"""Student model (spec §4, T2.2): Marco-Nano (HF Qwen3-MoE) with ternary routed experts. Each tensor class follows a precision map (configs/precision/.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/.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, `.sub_scales` (int8 m) beside `.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()