| |
| """D-57 · the teacher's routed experts quantised by our own trainer grids (round-to-nearest with the trainer's scale |
| search, no training) and written as a GGUF: the QAT starting point, measured in llama.cpp beside its own PTQ files. |
| |
| python scripts/rtn_gguf.py --method iq4xs --out teacher-ours-iq4xs-rtn.gguf |
| python scripts/rtn_gguf.py --ckpt runs/<run>/full_N/state.pt --method int2 --down-method iq4nl --out x.gguf |
| |
| With --ckpt, a trained student instead: its fp32 expert latents re-quantised (gate/up with --method, the down |
| projection with --down-method, default the method's own) and its trained attention and norms; the frozen router |
| and embedding are the teacher's. The checkpoint's torchao optimiser state is skipped (gen has no torchao). |
| |
| Experts go through `quant.ternary.export_packed` with `model.student.proj_spec` (so the down projection gets its |
| 32-block grid), everything else is written as the teacher holds it (bf16; norms and the router f32), in the template |
| GGUF's tensor order; `export.gguf.check_round_trip` then confirms the GGUF dequantises to the packed weights. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import sys |
| from pathlib import Path |
|
|
| import torch |
| from safetensors import safe_open |
|
|
| REPO = Path(__file__).resolve().parent.parent |
| sys.path.insert(0, str(REPO / "src")) |
| from lowbit.export.gguf import check_round_trip, gguf_module, student_tensors, write_gguf |
| from lowbit.model.student import proj_spec |
| from lowbit.quant.importance import load as load_importance |
| from lowbit.quant.ternary import GRIDS, METHODS, QuantSpec, export_packed |
|
|
| TEACHER = Path("/path/to/data/models/Marco-Nano-Instruct") |
| TEMPLATE = Path("/path/to/data/ptq/teacher-bf16.gguf") |
| LLAMA = Path("/path/to/llama.cpp") |
|
|
|
|
| def load_ckpt(path: str) -> dict: |
| """torch.load of a full checkpoint with torchao's optimiser classes replaced by inert stubs.""" |
| import pickle |
| import types |
|
|
| class StubT(torch.Tensor): |
| @classmethod |
| def __torch_dispatch__(cls, func, types_, args=(), kwargs=None): |
| raise RuntimeError("optimiser-state stub") |
|
|
| def __setstate__(self, st): |
| pass |
|
|
| class StubO: |
| def __init__(self, *a, **k): |
| pass |
|
|
| def __setstate__(self, st): |
| pass |
|
|
| class U(pickle.Unpickler): |
| def find_class(self, mod, name): |
| if mod.startswith("torchao"): |
| return StubT if ("Tensor" in name or "State" in name or "Optim" in name) else StubO |
| return super().find_class(mod, name) |
|
|
| pm = types.SimpleNamespace(Unpickler=U, load=pickle.load, __name__="pickle") |
| return torch.load(path, map_location="cpu", mmap=True, weights_only=False, pickle_module=pm) |
|
|
|
|
| def main() -> None: |
| p = argparse.ArgumentParser() |
| p.add_argument("--method", required=True, choices=METHODS) |
| p.add_argument("--out", required=True) |
| p.add_argument("--ckpt", help="a full checkpoint's state.pt: re-quantise its trained latents") |
| p.add_argument("--down-method", choices=METHODS) |
| p.add_argument("--importance", help="D-57 quant.importance safetensors: weight the scale search") |
| p.add_argument("--dense-as", choices=["bf16", "q8_0"], default="bf16", |
| help="attention and embedding: bf16, or q8_0 (the phone build)") |
| a = p.parse_args() |
| gguf = gguf_module(LLAMA) |
| def mk(m: str) -> QuantSpec: |
| return QuantSpec(method=m, group_size=GRIDS[m].group if m in GRIDS else 128) |
| spec = mk(a.method) |
| down_spec = mk(a.down_method) if a.down_method else proj_spec(spec, "down") |
| params = load_ckpt(a.ckpt)["params"] if a.ckpt else {} |
| imp = load_importance(a.importance) if a.importance else {} |
| cfg = json.loads((TEACHER / "config.json").read_text()) |
| wmap = json.loads((TEACHER / "model.safetensors.index.json").read_text())["weight_map"] |
| files = {f: safe_open(TEACHER / f, "pt") for f in set(wmap.values())} |
|
|
| def get(name: str) -> torch.Tensor: |
| return files[wmap[name]].get_tensor(name) |
|
|
| dense, packed = {}, {} |
| for name in wmap: |
| if ".mlp.experts." not in name: |
| dense[name] = params[name] if name in params else get(name) |
| for i in range(cfg["num_hidden_layers"]): |
| pre = f"model.layers.{i}.mlp.experts" |
| n = cfg["num_experts"] |
| if params: |
| gate_up, down = params[f"{pre}.gate_up_latent"], params[f"{pre}.down_latent"] |
| else: |
| gate = torch.stack([get(f"{pre}.{e}.gate_proj.weight") for e in range(n)]) |
| up = torch.stack([get(f"{pre}.{e}.up_proj.weight") for e in range(n)]) |
| down = torch.stack([get(f"{pre}.{e}.down_proj.weight") for e in range(n)]) |
| gate_up = torch.cat([gate, up], 1) |
| packed[f"{pre}.gate_up_proj"] = export_packed(gate_up.float(), proj_spec(spec, "gate_up"), |
| importance=imp.get(f"layers.{i}.gate_up")) |
| packed[f"{pre}.down_proj"] = export_packed(down.float(), down_spec, importance=imp.get(f"layers.{i}.down")) |
| print(f"layer {i} quantised", flush=True) |
| info = write_gguf(TEMPLATE, a.out, student_tensors(dense, packed, gguf, dense_as=a.dense_as), gguf) |
| rt = check_round_trip(a.out, packed, gguf) |
| print(json.dumps({**info, "bit_exact": rt["bit_exact"], "mismatched": rt["mismatched_tensors"][:3]}, indent=1)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
