code/src/lowbit/model/student.py

"""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 = ()  # A2: attention projections made ternary (e.g. q_proj, o_proj)
    int8_layers: tuple = ()   # A4: expert layers kept 8-bit per output channel instead of ternary
    importance: str | None = None  # D-57: quant.importance safetensors weighting the N-bit scale search
    codebook: str | None = None    # D-58: a `scripts/cb_import.py` directory: the experts start as that llama.cpp file
    codebook_mode: str = "frozen"  # D-58: "frozen" codes (scales train) or "free" (fp32 latents choose the entries)

    @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"))


# ---------------------------------------------------------------- ternary experts

# D-51/D-57: the down projection's rows (moe_intermediate_size, 384 in Marco) cannot take 256-wide super-blocks, so
# these methods train it on a 32-block grid: int3 on int4 (Q4_0), int3nl and iq4xs on iq4nl (IQ4_NL).
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:  # D-58: frozen codes, trained scales
        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"):  # free codes
            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)  # the latents hold the teacher's bf16 values exactly
    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  # autograd keeps what backward needs; the module keeps nothing
    mod.down_proj = None


def _make_ternary(experts: torch.nn.Module, spec: QuantSpec, state: _State) -> None:
    experts._compute_dtype = experts.gate_up_proj.dtype  # bf16 for Marco; the CPU tests run tiny models in fp32
    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"  # `train.optim`: optim.lr.scale, else optim.lr.ternary
        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"  # it moves in integer units
            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)  # exact: llama.cpp's fp32 weights
            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)."""
    # D-46: every layer shares one compiled frame; 4,096- and 8,192-token batches, the checkpoint's no-grad forward
    # and its recompute, and the in-process eval each add an entry. Past the default 8 a frame silently runs eager.
    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:  # D-57: non-persistent buffers, on the latents' device, read by the quantiser
            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"):  # A2's ternary attention: the ternary group
            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)


# ---------------------------------------------------------------- mask modes

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)  # position of each token's document start
    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:  # dense 4-D boolean mask (True = attend), the reference path
        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 students

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-58: the source file's blocks with the trained d and s written in
                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"):  # free codes: encode the entries the forward chooses
                    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):  # A4: int8 values and fp32 per-channel scales, as the forward used them
                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:  # A2
            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"):  # router fp32 if fp32; the rest bf16
            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:  # D-58: raw blocks; their v is in the codebook init (`dequant_codebook`)
            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:  # D-51: N-bit
            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():  # A2's ternary attention
        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", {}):  # A4's int8 expert layers
        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  # non-persistent buffers are not in the state dict: rebuild them off meta
    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()