| """Reading checkpoints and adapters, and reading/writing llama.cpp control-vector GGUFs. |
| |
| Tensors are read straight from safetensors by byte offset, so a 60 GB checkpoint is never loaded whole: |
| the fit reads one layer at a time. bf16 is upcast to fp32 exactly (the bf16 bits are the high half of |
| the fp32 word). |
| """ |
| from __future__ import annotations |
|
|
| import json |
| import struct |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| PREFIXES = ("model.language_model.", "model.") |
| N_LAYERS = 60 |
|
|
|
|
| def _header(path: Path) -> tuple[int, dict]: |
| with open(path, "rb") as fh: |
| n = struct.unpack("<Q", fh.read(8))[0] |
| return n, json.loads(fh.read(n)) |
|
|
|
|
| _HEADERS: dict[str, tuple[int, dict]] = {} |
|
|
|
|
| def read_tensor(path: Path, name: str) -> np.ndarray: |
| """One tensor from a safetensors file as fp32.""" |
| key = str(path) |
| if key not in _HEADERS: |
| _HEADERS[key] = _header(path) |
| n, hd = _HEADERS[key] |
| meta = hd[name] |
| a, b = meta["data_offsets"] |
| dtype = {"BF16": np.uint16, "F16": np.float16, "F32": np.float32}[meta["dtype"]] |
| raw = np.fromfile(path, dtype=dtype, count=(b - a) // np.dtype(dtype).itemsize, offset=8 + n + a) |
| raw = raw.reshape(meta["shape"]) |
| if meta["dtype"] == "BF16": |
| return (raw.astype(np.uint32) << 16).view(np.float32) |
| return raw.astype(np.float32) |
|
|
|
|
| def raw_bytes_equal(path_a: Path, name_a: str, path_b: Path, name_b: str, chunk: int = 1 << 26) -> bool: |
| """Whether two stored tensors are bit-identical (same dtype, shape and bytes), read in chunks.""" |
| na, ha = _header(path_a) |
| nb, hb = _header(path_b) |
| ma, mb = ha[name_a], hb[name_b] |
| if ma["dtype"] != mb["dtype"] or ma["shape"] != mb["shape"]: |
| return False |
| (a0, a1), (b0, _) = ma["data_offsets"], mb["data_offsets"] |
| with open(path_a, "rb") as fa, open(path_b, "rb") as fb: |
| fa.seek(8 + na + a0) |
| fb.seek(8 + nb + b0) |
| left = a1 - a0 |
| while left: |
| k = min(chunk, left) |
| if fa.read(k) != fb.read(k): |
| return False |
| left -= k |
| return True |
|
|
|
|
| class Checkpoint: |
| """Tensor access by name over one safetensors checkpoint directory (sharded or not).""" |
|
|
| def __init__(self, directory: str | Path): |
| self.dir = Path(directory) |
| index = self.dir / "model.safetensors.index.json" |
| if index.exists(): |
| self.map = json.loads(index.read_text())["weight_map"] |
| else: |
| single = self.dir / "model.safetensors" |
| self.map = {k: single.name for k in _header(single)[1] if k != "__metadata__"} |
| prefix = next((p for p in PREFIXES if any(k.startswith(p + "layers.") for k in self.map)), None) |
| assert prefix, f"{self.dir}: no decoder layers found" |
| self.prefix: str = prefix |
| layers = {int(k[len(self.prefix) + 7:].split(".")[0]) for k in self.map if k.startswith(self.prefix + "layers.")} |
| |
| assert layers == set(range(N_LAYERS)), f"{self.dir}: expected {N_LAYERS} decoder layers, found {len(layers)}" |
|
|
| def file(self, name: str) -> Path: |
| return self.dir / self.map[self.prefix + name] |
|
|
| def has(self, name: str) -> bool: |
| return self.prefix + name in self.map |
|
|
| def get(self, name: str) -> np.ndarray: |
| """name relative to the language-model prefix, e.g. 'layers.3.mlp.up_proj.weight'.""" |
| return read_tensor(self.file(name), self.prefix + name) |
|
|
| def layer(self, L: int) -> dict[str, np.ndarray]: |
| stem = f"{self.prefix}layers.{L}." |
| return {k[len(stem):]: read_tensor(self.dir / f, k) for k, f in self.map.items() if k.startswith(stem)} |
|
|
| def same_as(self, other: "Checkpoint", name: str) -> bool: |
| return raw_bytes_equal(self.file(name), self.prefix + name, other.file(name), other.prefix + name) |
|
|
| def embedding_rows(self, token_ids: np.ndarray) -> np.ndarray: |
| """Rows of embed_tokens for the given ids as fp32, read row by row (only the rows used).""" |
| name = self.prefix + "embed_tokens.weight" |
| path = self.file("embed_tokens.weight") |
| n, hd = _header(path) |
| meta = hd[name] |
| assert meta["dtype"] == "BF16", meta["dtype"] |
| width = meta["shape"][1] |
| start = 8 + n + meta["data_offsets"][0] |
| table = {} |
| with open(path, "rb") as fh: |
| for t in np.unique(token_ids): |
| fh.seek(start + int(t) * width * 2) |
| table[int(t)] = (np.frombuffer(fh.read(width * 2), np.uint16).astype(np.uint32) << 16).view(np.float32) |
| return np.stack([np.stack([table[int(t)] for t in row]) for row in token_ids]) |
|
|
|
|
| class Adapter: |
| """A PEFT LoRA adapter directory (adapter_config.json + adapter_model.safetensors).""" |
|
|
| def __init__(self, directory: str | Path, base_prefix: str = "model.language_model."): |
| self.dir = Path(directory) |
| cfg = json.loads((self.dir / "adapter_config.json").read_text()) |
| assert not cfg.get("use_rslora") and not cfg.get("use_dora"), "rsLoRA / DoRA adapters are not handled" |
| self.rank = int(cfg["r"]) |
| self.scale = cfg["lora_alpha"] / self.rank |
| self.file = self.dir / "adapter_model.safetensors" |
| self.keys = set(_header(self.file)[1]) |
| self.stem = f"base_model.model.{base_prefix}layers." |
|
|
| def families(self, L: int) -> list[str]: |
| """Projection names this adapter trains in layer L, e.g. 'self_attn.q_proj'.""" |
| stem = f"{self.stem}{L}." |
| return sorted({k[len(stem):-len(".lora_A.weight")] for k in self.keys |
| if k.startswith(stem) and k.endswith(".lora_A.weight")}) |
|
|
| def pair(self, L: int, family: str) -> tuple[np.ndarray, np.ndarray]: |
| """(A, scale * B) for one projection: the update is (scale * B) @ A.""" |
| k = f"{self.stem}{L}.{family}" |
| A = read_tensor(self.file, f"{k}.lora_A.weight") |
| B = read_tensor(self.file, f"{k}.lora_B.weight") |
| return A, self.scale * B |
|
|
|
|
| |
| |
| |
|
|
| def write_gguf(path: str | Path, vectors: np.ndarray, name: str | None = None, alpha: float = 1.0) -> None: |
| """Write alpha * vectors[1:] as a control-vector GGUF, so it loads at strength 1.0.""" |
| from gguf import GGUFWriter |
|
|
| vectors = np.asarray(vectors, np.float32) |
| assert vectors.shape[0] == N_LAYERS, vectors.shape |
| w = GGUFWriter(str(path), arch="controlvector") |
| w.add_string("controlvector.model_hint", "gemma4") |
| w.add_int32("controlvector.layer_count", N_LAYERS - 1) |
| w.add_float32("controlvector.alpha", float(alpha)) |
| if name: |
| w.add_string("general.name", name) |
| for L in range(1, N_LAYERS): |
| w.add_tensor(f"direction.{L}", np.ascontiguousarray(alpha * vectors[L], np.float32)) |
| w.write_header_to_file() |
| w.write_kv_data_to_file() |
| w.write_tensors_to_file() |
| w.close() |
|
|
|
|
| def read_gguf(path: str | Path) -> np.ndarray: |
| """A control-vector GGUF's direction.L tensors as an (N_LAYERS, d) array; layer 0 stays zero.""" |
| from gguf import GGUFReader |
|
|
| r = GGUFReader(str(path)) |
| t = {x.name: np.asarray(x.data, np.float32) for x in r.tensors if x.name.startswith("direction.")} |
| assert set(t) == {f"direction.{L}" for L in range(1, N_LAYERS)}, (path, len(t)) |
| out = np.zeros((N_LAYERS, len(t["direction.1"])), np.float32) |
| for k, v in t.items(): |
| out[int(k.split(".")[1])] = v |
| return out |
|
|
