| """Student model (spec §4, T2.2): Marco-Nano (HF Qwen3-MoE) with ternary routed experts. |
| |
| Each tensor class follows a precision map (configs/precision/<arm>.yaml): |
| - experts: fp32 latents are the parameters (`gate_up_latent` [E, 2I, H], `down_latent` [E, H, I]). A forward |
| pre-hook builds the layer's bf16 weights just in time and hands them to HF's experts implementation (the |
| eager loop, grouped_mm, ...; T2.3 picks): ternary fake-quant with the STE when quantisation is on, the |
| exact bf16 teacher weights when it is off. A post-hook drops them again, so only the running layer's |
| bf16 expert weights exist; with activation checkpointing the backward rebuilds them. |
| - attention and norms: bf16 working weights, exactly the teacher's; the trainer keeps fp32 masters (T2.5). |
| The map's `fp32` for norms is their master precision: bf16 activations need bf16 norm weights. |
| - router: computed exactly as stock Marco computes it (bf16 matmul, fp32 softmax, top-k renormalised), so the |
| quant-off student is bitwise the teacher and routing starts identical (D-28: an fp32 router changed top-1 on |
| 3% of positions). `form: fp32` in a precision map still selects an fp32 router (fp32 logits); trained in A1 |
| with an fp32 master in the trainer. |
| - embedding / LM head: tied, bf16, frozen. |
| |
| D-51's N-bit experts (int2 / int3 / int4, `quant.ternary.GRIDS`) use the same latents and hooks; only the quantiser |
| differs. int3's down projection is int4 (`proj_spec`): Q3_K needs rows of a multiple of 256, the down rows are 384. |
| |
| With quantisation off the student is the teacher: same weights, same kernels, bitwise-equal outputs (verified on |
| real Marco, T2.2). The teacher cache writer (T1.6) runs this module with quantisation off. |
| |
| Mask modes (spec §5): `causal_packed` (plain causal attention over the 4,096-token sequence) and |
| `doc_masked` (causal within each document, position ids reset per document). |
| """ |
| from __future__ import annotations |
|
|
| import hashlib |
| import json |
| from dataclasses import dataclass |
| from pathlib import Path |
|
|
| import torch |
| import yaml |
|
|
| from lowbit.quant import codebook as cbk |
| from lowbit.quant.ternary import GRIDS, QuantSpec, dequant_packed, export_packed, fake_quant, init_scale |
|
|
| MASK_MODES = ("causal_packed", "doc_masked") |
| TENSOR_CLASSES = ("experts", "attention", "norms", "router", "embedding") |
|
|
|
|
| @dataclass |
| class Precision: |
| arm: str |
| experts: QuantSpec |
| trainable: dict[str, bool] |
| router_fp32: bool = True |
| lr_scale: dict | None = None |
| attn_ternary: tuple = () |
| int8_layers: tuple = () |
| importance: str | None = None |
| codebook: str | None = None |
| codebook_mode: str = "frozen" |
|
|
| @classmethod |
| def load(cls, path: str | Path) -> "Precision": |
| """configs/precision/<arm>.yaml; experts default to Q-A absmean, g128, plain STE (spec §4).""" |
| c = yaml.safe_load(open(path)) |
| e = c["experts"] |
| spec = QuantSpec(method=e.get("method", "absmean"), group_size=e.get("group", 128), ste=e.get("ste", "plain"), |
| clip_threshold=e.get("clip_threshold", 1.5)) |
| trainable = {k: bool(c[k]["trainable"]) for k in TENSOR_CLASSES} |
| lr = {k: c[k]["lr_scale"] for k in TENSOR_CLASSES if "lr_scale" in c[k]} |
| return cls(arm=c.get("arm", Path(path).stem), experts=spec, trainable=trainable, |
| router_fp32=c["router"].get("form") == "fp32", lr_scale=lr or None, |
| attn_ternary=tuple(c["attention"].get("ternary") or ()), |
| int8_layers=tuple(e.get("int8_per_channel_layers") or ()), |
| importance=e.get("importance"), codebook=e.get("codebook"), |
| codebook_mode=e.get("codebook_mode", "frozen")) |
|
|
|
|
| |
|
|
| |
| |
| DOWN_METHOD = {"int3": "int4", "int3nl": "iq4nl", "iq4xs": "iq4nl", "int2nl": "iq4nl"} |
|
|
|
|
| def proj_spec(spec: QuantSpec, name: str) -> QuantSpec: |
| """The quantiser of one expert projection ("gate_up" or "down"); see DOWN_METHOD.""" |
| if name == "down" and spec.method in DOWN_METHOD: |
| m = DOWN_METHOD[spec.method] |
| return QuantSpec(method=m, group_size=GRIDS[m].group, ste=spec.ste) |
| return spec |
|
|
|
|
| class _State: |
| """Shared switch: every experts module of one student reads it.""" |
|
|
| def __init__(self, quant: bool): |
| self.quant = quant |
|
|
|
|
| def fake_quant_int8(w: torch.Tensor) -> torch.Tensor: |
| """Symmetric 8-bit per output channel (last dim = inputs), straight-through gradient (A4, student only).""" |
| s = w.detach().abs().amax(-1, keepdim=True).clamp_min(1e-12) / 127 |
| w_hat = (w.detach() / s).round().clamp(-127, 127) * s |
| return w + (w_hat - w).detach() |
|
|
|
|
| def _expert_weight(mod: torch.nn.Module, name: str) -> torch.Tensor: |
| fmt = getattr(mod, "_cb", {}).get(name) |
| if fmt is not None: |
| if not mod._ternary_state.quant: |
| raise RuntimeError("a codebook student has no unquantised weights (its teacher is a separate model)") |
| if hasattr(mod, f"{name}_latent"): |
| return cbk.weight_free(fmt, getattr(mod, f"{name}_latent"), getattr(mod, f"{name}_d0"), |
| getattr(mod, f"{name}_dlog"), getattr(mod, f"{name}_slat", None), |
| getattr(mod, f"{name}_grid")).to(mod._compute_dtype) |
| return cbk.weight(fmt, getattr(mod, f"{name}_v"), getattr(mod, f"{name}_d0"), getattr(mod, f"{name}_dlog"), |
| getattr(mod, f"{name}_slat", None)).to(mod._compute_dtype) |
| latent = getattr(mod, f"{name}_latent") |
| if not mod._ternary_state.quant: |
| return latent.to(mod._compute_dtype) |
| if getattr(mod, "_int8", False): |
| return fake_quant_int8(latent).to(mod._compute_dtype) |
| scale = getattr(mod, f"{name}_scale", None) |
| return fake_quant(latent, mod._specs[name], scale, getattr(mod, f"{name}_imp", None)).to(mod._compute_dtype) |
|
|
|
|
| def _pre_hook(mod, args): |
| mod.gate_up_proj = _expert_weight(mod, "gate_up") |
| mod.down_proj = _expert_weight(mod, "down") |
|
|
|
|
| def _post_hook(mod, args, output): |
| mod.gate_up_proj = None |
| mod.down_proj = None |
|
|
|
|
| def _make_ternary(experts: torch.nn.Module, spec: QuantSpec, state: _State) -> None: |
| experts._compute_dtype = experts.gate_up_proj.dtype |
| for name in ("gate_up", "down"): |
| w = experts._parameters.pop(f"{name}_proj") |
| latent = torch.nn.Parameter(w.detach().float()) |
| experts.register_parameter(f"{name}_latent", latent) |
| if spec.method == "lsq": |
| experts.register_parameter(f"{name}_scale", torch.nn.Parameter(init_scale(latent, spec))) |
| setattr(experts, f"{name}_proj", None) |
| experts._ternary_spec = spec |
| experts._specs = {name: proj_spec(spec, name) for name in ("gate_up", "down")} |
| experts._ternary_state = state |
| experts.register_forward_pre_hook(_pre_hook) |
| experts.register_forward_hook(_post_hook) |
|
|
|
|
| class CodebookInit: |
| """A `scripts/cb_import.py` directory, read one tensor at a time.""" |
|
|
| def __init__(self, path: str | Path): |
| from safetensors import safe_open |
| self.path = Path(path) |
| self.formats = json.loads((self.path / "manifest.json").read_text())["formats"] |
| self._f = safe_open(str(self.path / "experts.safetensors"), framework="pt") |
|
|
| def cat(self, i: int, parts: tuple, field: str) -> torch.Tensor: |
| """Layer i's `field` of the given projections, concatenated along the rows (gate first, as HF fuses them).""" |
| return torch.cat([self._f.get_tensor(f"layers.{i}.{p}.{field}") for p in parts], 1) |
|
|
| def grid(self, name: str) -> torch.Tensor: |
| from safetensors.torch import load_file |
| return load_file(str(self.path / "grids.safetensors"))[name] |
|
|
| def fmt(self, i: int, parts: tuple) -> cbk.Format: |
| names = {self.formats[f"layers.{i}.{p}"] for p in parts} |
| if len(names) != 1: |
| raise ValueError(f"layer {i}: {parts} have different formats {names}; the fused gate_up needs one") |
| return cbk.FORMATS[names.pop()] |
|
|
|
|
| PARTS = {"gate_up": ("gate", "up"), "down": ("down",)} |
|
|
|
|
| def _make_codebook(experts: torch.nn.Module, i: int, init: CodebookInit, state: _State, mode: str = "frozen") -> None: |
| """D-58: layer i's experts as a llama.cpp file's codes. v (int8) and d0 are non-persistent buffers; `dlog` (per |
| block) and `slat` (per sub-block, the integer scale as a float) are the parameters (`quant.codebook`). "free" |
| mode replaces v with an fp32 `latent` at the file's dequantised weights, and the codebook as a buffer.""" |
| dev = experts.gate_up_proj.device |
| experts._compute_dtype = experts.gate_up_proj.dtype |
| experts._cb = {} |
| for name, parts in PARTS.items(): |
| experts._parameters.pop(f"{name}_proj") |
| fmt = init.fmt(i, parts) |
| v = init.cat(i, parts, "v").to(dev) |
| d0 = init.cat(i, parts, "d").float().to(dev) |
| experts.register_buffer(f"{name}_d0", d0, persistent=False) |
| dlog = torch.nn.Parameter(torch.zeros_like(d0)) |
| dlog._mt_lr = "scale" |
| experts.register_parameter(f"{name}_dlog", dlog) |
| s0 = None |
| if fmt.name != "IQ4_NL": |
| s0 = init.cat(i, parts, "s").float().to(dev) |
| slat = torch.nn.Parameter(s0.clone()) |
| slat._mt_lr = "subscale" |
| experts.register_parameter(f"{name}_slat", slat) |
| if mode == "free": |
| e = cbk.eff_torch(fmt, d0, s0).unsqueeze(-1) |
| lat = (v.float().unflatten(-1, (-1, fmt.sub)) * e).flatten(-2) |
| experts.register_parameter(f"{name}_latent", torch.nn.Parameter(lat)) |
| experts.register_buffer(f"{name}_grid", init.grid(fmt.name).to(dev), persistent=False) |
| del v |
| else: |
| experts.register_buffer(f"{name}_v", v, persistent=False) |
| setattr(experts, f"{name}_proj", None) |
| experts._cb[name] = fmt |
| experts._ternary_state = state |
| experts.register_forward_pre_hook(_pre_hook) |
| experts.register_forward_hook(_post_hook) |
|
|
|
|
| def _router_fp32_forward(self, hidden_states): |
| h = hidden_states.reshape(-1, self.hidden_dim) |
| router_logits = torch.nn.functional.linear(h.float(), self.weight.float()) |
| probs = torch.softmax(router_logits, dim=-1) |
| top_v, top_i = torch.topk(probs, self.top_k, dim=-1) |
| if self.norm_topk_prob: |
| top_v = top_v / top_v.sum(dim=-1, keepdim=True) |
| return router_logits, top_v.to(hidden_states.dtype), top_i |
|
|
|
|
| def enable_checkpointing(model) -> None: |
| """Per-layer activation checkpointing, non-reentrant and without early stop. An early-stopped recompute ends |
| inside the experts' forward, before `_post_hook`, so every layer kept its recomputed ternary weights and their |
| codes after its backward: +1.65 GB per Marco layer, 46 GB in all (T2.3).""" |
| model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False, "early_stop": False}) |
|
|
|
|
| def compile_layers(model) -> None: |
| """torch.compile each decoder layer, inside its checkpoint. The quantiser compiles with it: its codes are exact |
| (fp64 group sums, comparisons instead of division), so compiled and eager runs quantise identically (T2.3).""" |
| |
| |
| torch._dynamo.config.recompile_limit = max(torch._dynamo.config.recompile_limit, 32) |
| for layer in model.model.layers: |
| layer.compile(dynamic=False) |
|
|
|
|
| def _linear_pre_hook(mod, args): |
| lat = mod.weight_latent |
| if not mod._ternary_state.quant: |
| mod.weight = lat.to(mod._compute_dtype) |
| else: |
| mod.weight = fake_quant(lat, mod._ternary_spec).to(mod._compute_dtype) |
|
|
|
|
| def _linear_post_hook(mod, args, output): |
| mod.weight = None |
|
|
|
|
| def _make_ternary_linear(lin: torch.nn.Linear, spec: QuantSpec, state: _State) -> None: |
| """A2: a Linear whose weight is a ternary fake-quant of an fp32 latent, built just in time like the experts'.""" |
| w = lin._parameters.pop("weight") |
| lin._compute_dtype = w.dtype |
| lin.register_parameter("weight_latent", torch.nn.Parameter(w.detach().float())) |
| lin.weight = None |
| lin._ternary_spec, lin._ternary_state = spec, state |
| lin.register_forward_pre_hook(_linear_pre_hook) |
| lin.register_forward_hook(_linear_post_hook) |
|
|
|
|
| def build_student(model, precision: Precision, quant: bool = True): |
| """Turn a loaded bf16 HF Qwen3-MoE model into the student in place; returns it.""" |
| import types |
| if precision.attn_ternary and precision.experts.method in GRIDS: |
| raise NotImplementedError("A2's quantised attention is built for the ternary methods only") |
| state = _State(quant) |
| if precision.codebook and (precision.attn_ternary or precision.int8_layers or precision.importance): |
| raise NotImplementedError("D-58 codebook students take no A2/A4/D-57 options") |
| cb_init = CodebookInit(precision.codebook) if precision.codebook else None |
| imp = None |
| if precision.importance: |
| from lowbit.quant.importance import load as load_importance |
| imp = load_importance(precision.importance) |
| for i, layer in enumerate(model.model.layers): |
| if cb_init is not None: |
| _make_codebook(layer.mlp.experts, i, cb_init, state, precision.codebook_mode) |
| else: |
| _make_ternary(layer.mlp.experts, precision.experts, state) |
| layer.mlp.experts._int8 = i in precision.int8_layers |
| if imp is not None: |
| ex = layer.mlp.experts |
| for name in ("gate_up", "down"): |
| ex.register_buffer(f"{name}_imp", imp[f"layers.{i}.{name}"].to(ex.gate_up_latent.device), |
| persistent=False) |
| for name in precision.attn_ternary: |
| _make_ternary_linear(getattr(layer.self_attn, name), precision.experts, state) |
| if precision.router_fp32: |
| gate = layer.mlp.gate |
| gate.weight = torch.nn.Parameter(gate.weight.detach().float()) |
| gate.forward = types.MethodType(_router_fp32_forward, gate) |
| model._ternary_state = state |
| model._precision = precision |
| for p in model.parameters(): |
| p.requires_grad_(False) |
| for cls, params in param_classes(model).items(): |
| for p in params: |
| p.requires_grad_(precision.trainable[cls]) |
| return model |
|
|
|
|
| def set_quant(model, on: bool) -> None: |
| model._ternary_state.quant = on |
|
|
|
|
| def param_classes(model) -> dict[str, list[torch.nn.Parameter]]: |
| """Parameters by tensor class (spec §4), for trainability and the trainer's LR groups.""" |
| out: dict[str, list] = {k: [] for k in TENSOR_CLASSES} |
| for name, p in model.named_parameters(): |
| if ".experts." in name or name.endswith(".weight_latent"): |
| out["experts"].append(p) |
| elif ".self_attn." in name and name.endswith("_proj.weight"): |
| out["attention"].append(p) |
| elif "norm" in name: |
| out["norms"].append(p) |
| elif ".mlp.gate." in name: |
| out["router"].append(p) |
| elif "embed_tokens" in name or "lm_head" in name: |
| out["embedding"].append(p) |
| else: |
| raise ValueError(f"unclassified parameter {name}") |
| return out |
|
|
|
|
| def load_student(model_dir: str | Path, precision: Precision | str | Path, quant: bool = True, device=None, |
| attn_impl: str = "sdpa", experts_impl: str | None = None): |
| from transformers import AutoModelForCausalLM |
| if not isinstance(precision, Precision): |
| precision = Precision.load(precision) |
| kw = {"experts_implementation": experts_impl} if experts_impl else {} |
| model = AutoModelForCausalLM.from_pretrained(model_dir, dtype=torch.bfloat16, attn_implementation=attn_impl, **kw) |
| if device is not None: |
| model.to(device) |
| return build_student(model, precision, quant) |
|
|
|
|
| |
|
|
| def doc_ids(doc_starts: list[list[int]], seq_len: int, device=None) -> torch.Tensor: |
| """[B, T] document index per position from each sequence's document start offsets.""" |
| ids = torch.zeros(len(doc_starts), seq_len, dtype=torch.long, device=device) |
| for b, starts in enumerate(doc_starts): |
| for s in starts: |
| if 0 < s < seq_len: |
| ids[b, s:] += 1 |
| return ids |
|
|
|
|
| def mask_inputs(doc_starts: list[list[int]] | None, seq_len: int, mode: str, attn_impl: str, device=None) -> dict: |
| """position_ids and attention_mask for model.forward under a mask mode.""" |
| if mode not in MASK_MODES: |
| raise ValueError(f"mask mode {mode!r} not in {MASK_MODES}") |
| if mode == "causal_packed" or doc_starts is None: |
| return {} |
| d = doc_ids(doc_starts, seq_len, device) |
| pos = torch.arange(seq_len, device=device).expand_as(d).clone() |
| first = torch.zeros_like(d) |
| for b, starts in enumerate(doc_starts): |
| for s in starts: |
| if 0 < s < seq_len: |
| first[b, s:] = s |
| pos = pos - first |
| if attn_impl == "flex_attention": |
| from torch.nn.attention.flex_attention import create_block_mask |
|
|
| def doc_causal(b, h, q, kv): |
| return (d[b, q] == d[b, kv]) & (q >= kv) |
| mask = create_block_mask(doc_causal, len(doc_starts), None, seq_len, seq_len, device=device) |
| else: |
| causal = torch.ones(seq_len, seq_len, dtype=torch.bool, device=device).tril() |
| mask = (d[:, :, None] == d[:, None, :]) & causal |
| mask = mask[:, None] |
| return {"position_ids": pos, "attention_mask": mask} |
|
|
|
|
| def final_hidden(model, input_ids: torch.Tensor, mode: str = "causal_packed", doc_starts=None) -> torch.Tensor: |
| """The final (normed) hidden states [B, T, H]; the losses apply the frozen LM head in chunks.""" |
| impl = model.config._attn_implementation |
| out = model.model(input_ids=input_ids, **mask_inputs(doc_starts, input_ids.shape[1], mode, impl, input_ids.device)) |
| return out.last_hidden_state |
|
|
|
|
| def logits(model, hidden: torch.Tensor) -> torch.Tensor: |
| return model.lm_head(hidden) |
|
|
|
|
| |
|
|
| PACKED_VERSION = 1 |
|
|
|
|
| def export_packed_student(model, out_dir: str | Path) -> dict: |
| """2-bit expert codes + fp16 scales, everything else bf16 (spec §4 "Export form"); ~2.7 GB for Marco. |
| D-51's N-bit experts add `expert_specs` (method and group per tensor) to the manifest and, for int3's Q3_K |
| tensors, `<key>.sub_scales` (int8 m) beside `<key>.scales` (d).""" |
| from safetensors.torch import save_file |
| out = Path(out_dir) |
| out.mkdir(parents=True, exist_ok=True) |
| spec = model._precision.experts |
| tensors, shapes, int8_shapes, linear_shapes, specs = {}, {}, {}, {}, {} |
| cb_init = CodebookInit(model._precision.codebook) if model._precision.codebook else None |
| cb_formats = {} |
| for i, layer in enumerate(model.model.layers): |
| ex = layer.mlp.experts |
| for name in ("gate_up", "down"): |
| key = f"model.layers.{i}.mlp.experts.{name}_proj" |
| fmt = getattr(ex, "_cb", {}).get(name) |
| if fmt is not None: |
| d, s = cbk.trained_fields(fmt, getattr(ex, f"{name}_d0").detach(), getattr(ex, f"{name}_dlog").detach(), |
| None if fmt.name == "IQ4_NL" else getattr(ex, f"{name}_slat").detach()) |
| if hasattr(ex, f"{name}_latent"): |
| lat = getattr(ex, f"{name}_latent").detach() |
| e = cbk.eff_torch(fmt, torch.from_numpy(d.astype("float32")).to(lat.device), |
| torch.from_numpy(s.astype("float32")).to(lat.device)).unsqueeze(-1) |
| idx, vals = cbk.free_codes(fmt, (lat.unflatten(-1, (-1, fmt.sub)) / e).flatten(-2), |
| getattr(ex, f"{name}_grid")) |
| tensors[f"{key}.blocks"] = torch.from_numpy(cbk.encode(fmt, idx, vals, d, s)) |
| shapes[key] = list(lat.shape) |
| else: |
| blocks = cb_init.cat(i, PARTS[name], "blocks").numpy() |
| tensors[f"{key}.blocks"] = torch.from_numpy(cbk.write_blocks(fmt, blocks, d, s)) |
| shapes[key] = list(getattr(ex, f"{name}_v").shape) |
| cb_formats[key] = fmt.name |
| continue |
| lat = getattr(ex, f"{name}_latent").detach() |
| if getattr(ex, "_int8", False): |
| s = lat.abs().amax(-1, keepdim=True).clamp_min(1e-12) / 127 |
| tensors[f"{key}.int8"] = (lat / s).round().clamp(-127, 127).to(torch.int8).contiguous() |
| tensors[f"{key}.int8_scales"] = s.float().contiguous() |
| int8_shapes[key] = list(lat.shape) |
| continue |
| p = export_packed(lat, ex._specs[name], getattr(ex, f"{name}_scale", None), |
| getattr(ex, f"{name}_imp", None)) |
| tensors[f"{key}.codes"] = p["codes"].contiguous() |
| tensors[f"{key}.scales"] = p["scales"].contiguous() |
| if "sub_scales" in p: |
| tensors[f"{key}.sub_scales"] = p["sub_scales"].contiguous() |
| if "method" in p: |
| specs[key] = [p["method"], p["group_size"]] |
| shapes[key] = list(p["shape"]) |
| for name in model._precision.attn_ternary: |
| lat = getattr(layer.self_attn, name).weight_latent.detach() |
| p = export_packed(lat, spec) |
| key = f"model.layers.{i}.self_attn.{name}.weight" |
| tensors[f"{key}.codes"] = p["codes"].contiguous() |
| tensors[f"{key}.scales"] = p["scales"].contiguous() |
| linear_shapes[key] = list(p["shape"]) |
| for name, p in model.named_parameters(): |
| if ".experts." not in name and not name.endswith(".weight_latent"): |
| keep = ".mlp.gate." in name and model._precision.router_fp32 |
| tensors[name] = p.detach().to(p.dtype if keep else torch.bfloat16).contiguous() |
| save_file({k: v.cpu() for k, v in tensors.items()}, str(out / "model.safetensors")) |
| model.config.save_pretrained(out) |
| manifest = {"packed_version": PACKED_VERSION, "arm": model._precision.arm, "group_size": spec.group_size, |
| "method": spec.method, "router_fp32": model._precision.router_fp32, "expert_shapes": shapes, |
| "int8_shapes": int8_shapes, "linear_shapes": linear_shapes, **({"expert_specs": specs} if specs else {}), |
| **({"codebook_formats": cb_formats, "codebook_init": str(cb_init.path)} if cb_formats else {}), |
| "sha256": hashlib.sha256((out / "model.safetensors").read_bytes()).hexdigest()} |
| (out / "packed_manifest.json").write_text(json.dumps(manifest, indent=1)) |
| return manifest |
|
|
|
|
| def read_packed(packed_dir: str | Path) -> tuple[dict, dict, dict]: |
| """(the plain tensors by HF name, the packed expert tensors by key as `dequant_packed` takes them, the manifest).""" |
| from safetensors.torch import load_file |
| d = Path(packed_dir) |
| man = json.loads((d / "packed_manifest.json").read_text()) |
| sd = load_file(str(d / "model.safetensors")) |
| specs = man.get("expert_specs", {}) |
| cbf = man.get("codebook_formats", {}) |
| packed = {} |
| for key, shape in man["expert_shapes"].items(): |
| if key in cbf: |
| packed[key] = {"blocks": sd.pop(f"{key}.blocks"), "format": cbf[key], "shape": tuple(shape), |
| "init": man["codebook_init"]} |
| continue |
| p = {"codes": sd.pop(f"{key}.codes"), "scales": sd.pop(f"{key}.scales"), "group_size": man["group_size"], |
| "shape": tuple(shape)} |
| if key in specs: |
| p["method"], p["group_size"] = specs[key] |
| if f"{key}.sub_scales" in sd: |
| p["sub_scales"] = sd.pop(f"{key}.sub_scales") |
| packed[key] = p |
| return sd, packed, man |
|
|
|
|
| def load_codebook_export(model, packed_dir: str | Path) -> None: |
| """D-58: continue a frozen-code student from one of its own packed exports. Each block's d and s become |
| dlog = log(d / d0) (round_fp16(d0·exp(dlog)) gives d back exactly) and slat = s; the export's other tensors |
| (attention, norms; the frozen router and embedding are the teacher's anyway) are copied in. Its export is then |
| the source export, byte for byte.""" |
| import numpy as np |
| sd, packed, _ = read_packed(packed_dir) |
| named = dict(model.named_parameters()) |
| with torch.no_grad(): |
| for i, layer in enumerate(model.model.layers): |
| ex = layer.mlp.experts |
| for name in PARTS: |
| if hasattr(ex, f"{name}_latent"): |
| raise NotImplementedError("free-code students keep latents, which exports do not carry") |
| p = packed[f"model.layers.{i}.mlp.experts.{name}_proj"] |
| fmt = cbk.FORMATS[p["format"]] |
| flat = p["blocks"].numpy().reshape(-1, fmt.nbytes) |
| d0 = getattr(ex, f"{name}_d0") |
| d = torch.from_numpy(cbk.get_d(flat).astype(np.float32)).reshape(d0.shape).to(d0.device) |
| getattr(ex, f"{name}_dlog").copy_(torch.where(d0 != 0, torch.log(d / d0), torch.zeros_like(d0))) |
| if fmt.name != "IQ4_NL": |
| slat = getattr(ex, f"{name}_slat") |
| slat.copy_(torch.from_numpy(cbk.get_s(fmt, flat).astype(np.float32)).reshape(slat.shape)) |
| for n, t in sd.items(): |
| if n in named: |
| named[n].copy_(t.to(named[n].device, named[n].dtype)) |
|
|
|
|
| def dequant_codebook(key: str, p: dict, dtype: torch.dtype = torch.bfloat16) -> torch.Tensor: |
| """D-58: a packed codebook tensor's weights (its v from the init the student started from).""" |
| import re |
| i, name = re.fullmatch(r"model\.layers\.(\d+)\.mlp\.experts\.(gate_up|down)_proj", key).groups() |
| v = CodebookInit(p["init"]).cat(int(i), PARTS[name], "v").numpy() |
| return torch.from_numpy(cbk.dequant(cbk.FORMATS[p["format"]], p["blocks"].numpy(), v)).to(dtype) |
|
|
|
|
| def load_packed_student(packed_dir: str | Path, device=None, attn_impl: str = "sdpa", experts_impl: str | None = None): |
| """A plain bf16 HF model whose expert weights are the dequantised codes: the fake-quant student exactly.""" |
| from transformers import AutoConfig, AutoModelForCausalLM |
| d = Path(packed_dir) |
| sd, packed, man = read_packed(d) |
| cfg = AutoConfig.from_pretrained(d) |
| kw = {"experts_implementation": experts_impl} if experts_impl else {} |
| with torch.device("meta"): |
| model = AutoModelForCausalLM.from_config(cfg, dtype=torch.bfloat16, attn_implementation=attn_impl, **kw) |
| state = {key: dequant_codebook(key, p) if "blocks" in p else dequant_packed(p) for key, p in packed.items()} |
| for key, shape in man.get("linear_shapes", {}).items(): |
| state[key] = dequant_packed({"codes": sd.pop(f"{key}.codes"), "scales": sd.pop(f"{key}.scales"), |
| "group_size": man["group_size"], "shape": tuple(shape)}) |
| for key in man.get("int8_shapes", {}): |
| state[key] = (sd.pop(f"{key}.int8").float() * sd.pop(f"{key}.int8_scales")).to(torch.bfloat16) |
| state.update(sd) |
| if cfg.tie_word_embeddings and "lm_head.weight" not in state: |
| state["lm_head.weight"] = state["model.embed_tokens.weight"] |
| model.load_state_dict(state, strict=True, assign=True) |
| model.tie_weights() |
| rope = model.model.rotary_emb |
| model.model.rotary_emb = type(rope)(config=cfg) |
| if man.get("router_fp32"): |
| import types |
| for layer in model.model.layers: |
| layer.mlp.gate.forward = types.MethodType(_router_fp32_forward, layer.mlp.gate) |
| if device is not None: |
| model.to(device) |
| return model.eval() |
|
|
