code/tools/build_checkpoint_v3.py

#!/usr/bin/env python3
"""Assemble the per-expert hybrid NVFP4+AQLM GLM-5.2 checkpoint.

Sources:
  - /tmp/glm52-dl/nvfp4-full: all non-expert tensors verbatim; per-expert
    NVFP4 tensors for the whole-NVFP4 layers AND for hot experts of hybrid
    layers (compacted into fused arrays)
  - /data/glm52-aqlm-parts/layer_N.pt: AQLM codes for base/cold experts
    (sliced per assignment; cold w2 = book-0 slice of the 2-book codes)
  - /data/glm52-expert-assignment.json: per-layer hot/cold expert ids

Output: /data/glm52-v3 (swap into /data/glm52 after validation).
"""

import json
import os
import shutil

import torch
from safetensors import safe_open
from safetensors.torch import save_file

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
SRC = "/tmp/glm52-dl/nvfp4-full"
PARTS = "/data/glm52-aqlm-parts"
ASSIGN = "/data/glm52-expert-assignment.json"
DST = "/data/glm52-v3"
SHARD_BYTES = 4 << 30

plan = json.load(open(os.path.join(ROOT, "hybrid_plan.json")))
NVFP4_LAYERS = set(plan["nvfp4_layers"])           # whole-layer NVFP4 (incl MTP)
assignment = {int(k): v for k, v in json.load(open(ASSIGN)).items()}
HYBRID_LAYERS = sorted(assignment)

import re

_EXPERT_RE = re.compile(r"model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.")


def keep(name: str) -> bool:
    m = _EXPERT_RE.match(name)
    if m:
        return int(m.group(1)) in NVFP4_LAYERS
    return True


class ShardWriter:
    def __init__(self, dst: str):
        self.dst = dst
        self.cur: dict[str, torch.Tensor] = {}
        self.cur_bytes = 0
        self.n = 0
        self.weight_map: dict[str, str] = {}
        self.total = 0
        self.files: list[str] = []

    def add(self, name: str, tensor: torch.Tensor):
        nb = tensor.numel() * tensor.element_size()
        if self.cur_bytes + nb > SHARD_BYTES and self.cur:
            self.flush()
        self.cur[name] = tensor
        self.cur_bytes += nb
        self.total += nb

    def flush(self):
        if not self.cur:
            return
        self.n += 1
        fname = f"model-{self.n:05d}.safetensors"
        save_file(self.cur, os.path.join(self.dst, fname))
        for k in self.cur:
            self.weight_map[k] = fname
        print(f"  wrote {fname} ({self.cur_bytes/1e9:.2f} GB)", flush=True)
        self.files.append(fname)
        self.cur = {}
        self.cur_bytes = 0

    def finalize(self):
        self.flush()
        total_n = self.n
        wm = {}
        for i, fname in enumerate(self.files, 1):
            new = f"model-{i:05d}-of-{total_n:05d}.safetensors"
            os.rename(os.path.join(self.dst, fname), os.path.join(self.dst, new))
            for k, v in self.weight_map.items():
                if v == fname:
                    wm[k] = new
        json.dump(
            {"metadata": {"total_size": self.total}, "weight_map": wm},
            open(os.path.join(self.dst, "model.safetensors.index.json"), "w"),
            indent=0,
        )
        print(f"index: {len(wm)} tensors, {self.total/1e9:.1f} GB total")


class ShardReader:
    def __init__(self, repo_dir: str):
        idx = json.load(open(os.path.join(repo_dir, "model.safetensors.index.json")))
        self.weight_map = idx["weight_map"]
        self.repo_dir = repo_dir
        self._open = {}

    def get(self, name: str):
        shard = self.weight_map[name]
        if shard not in self._open:
            self._open[shard] = safe_open(
                os.path.join(self.repo_dir, shard), framework="pt"
            )
        return self._open[shard].get_tensor(name)


def build_hybrid_layer(li: int, reader: ShardReader, writer: ShardWriter):
    n_exp, inter, hidden = 256, 2048, 6144
    hot = sorted(assignment[li]["hot"])
    cold = sorted(assignment[li]["cold"])
    hot_s, cold_s = set(hot), set(cold)
    base = [e for e in range(n_exp) if e not in hot_s and e not in cold_s]
    b_all = sorted(set(base) | cold_s)  # AQLM w13 group, ascending

    kind = torch.ones(n_exp, dtype=torch.int8)
    for e in hot:
        kind[e] = 0
    for e in cold:
        kind[e] = 2

    part = torch.load(
        os.path.join(PARTS, f"layer_{li}.pt"), map_location="cpu",
        weights_only=True,
    )
    assert part["books"]["w2"] == 2, f"layer {li} part lacks 2-book w2"
    p = f"model.layers.{li}.mlp.experts"

    b_idx = torch.tensor(b_all, dtype=torch.long)
    base_idx = torch.tensor(base, dtype=torch.long)
    cold_idx = torch.tensor(cold, dtype=torch.long)

    writer.add(f"{p}.hyb_kind", kind)
    writer.add(f"{p}.w13_codes", part["w13_codes"][b_idx].contiguous())
    writer.add(f"{p}.w13_codebooks", part["w13_codebooks"])
    writer.add(f"{p}.w13_scales", part["w13_scales"][b_idx].contiguous())
    writer.add(f"{p}.w2m_codes", part["w2_codes"][base_idx].contiguous())
    writer.add(f"{p}.w2m_codebooks", part["w2_codebooks"])
    writer.add(f"{p}.w2m_scales", part["w2_scales"][base_idx].contiguous())
    # cold: book-0 slice of the 2-book residual encoding
    writer.add(f"{p}.w2c_codes", part["w2_codes"][cold_idx, :1].clone())
    writer.add(f"{p}.w2c_codebooks", part["w2_codebooks"][:1].clone())
    writer.add(f"{p}.w2c_scales", part["w2_scales"][cold_idx].contiguous())

    # hot experts: compact NVFP4 arrays from the source repo
    na = len(hot)
    w13_packed = torch.empty(na, 2 * inter, hidden // 2, dtype=torch.uint8)
    w13_bscale = torch.empty(na, 2 * inter, hidden // 16, dtype=torch.uint8)
    w13_scale2 = torch.empty(na, 2, dtype=torch.float32)
    w2_packed = torch.empty(na, hidden, inter // 2, dtype=torch.uint8)
    w2_bscale = torch.empty(na, hidden, inter // 16, dtype=torch.uint8)
    w2_scale2 = torch.empty(na, 1, dtype=torch.float32)
    for j, e in enumerate(hot):
        ep = f"model.layers.{li}.mlp.experts.{e}"
        g_w = reader.get(f"{ep}.gate_proj.weight")
        u_w = reader.get(f"{ep}.up_proj.weight")
        d_w = reader.get(f"{ep}.down_proj.weight")
        w13_packed[j, :inter] = g_w
        w13_packed[j, inter:] = u_w
        w2_packed[j] = d_w
        w13_bscale[j, :inter] = reader.get(
            f"{ep}.gate_proj.weight_scale").view(torch.uint8)
        w13_bscale[j, inter:] = reader.get(
            f"{ep}.up_proj.weight_scale").view(torch.uint8)
        w2_bscale[j] = reader.get(
            f"{ep}.down_proj.weight_scale").view(torch.uint8)
        w13_scale2[j, 0] = reader.get(f"{ep}.gate_proj.weight_scale_2").float()
        w13_scale2[j, 1] = reader.get(f"{ep}.up_proj.weight_scale_2").float()
        w2_scale2[j, 0] = reader.get(f"{ep}.down_proj.weight_scale_2").float()
    writer.add(f"{p}.nvfp4_w13_packed", w13_packed)
    writer.add(f"{p}.nvfp4_w13_bscale", w13_bscale)
    writer.add(f"{p}.nvfp4_w13_scale2", w13_scale2)
    writer.add(f"{p}.nvfp4_w2_packed", w2_packed)
    writer.add(f"{p}.nvfp4_w2_bscale", w2_bscale)
    writer.add(f"{p}.nvfp4_w2_scale2", w2_scale2)

    print(f"layer {li}: hot={na} base={len(base)} cold={len(cold)}", flush=True)
    return {"n_nvfp4": na, "n_base": len(base), "n_cold": len(cold)}


def main():
    os.makedirs(DST, exist_ok=True)
    reader = ShardReader(SRC)
    writer = ShardWriter(DST)

    # 1. non-expert tensors + whole-NVFP4 layers, streamed shard by shard
    shards = sorted(set(reader.weight_map.values()))
    for shard in shards:
        with safe_open(os.path.join(SRC, shard), framework="pt") as f:
            names = [n for n in f.keys() if keep(n)]
            if not names:
                continue
            for n in names:
                writer.add(n, f.get_tensor(n))
        print(f"{shard}: kept {len(names)}", flush=True)

    # 2. hybrid layers
    layer_books = {}
    for li in HYBRID_LAYERS:
        layer_books[str(li)] = build_hybrid_layer(li, reader, writer)

    writer.finalize()

    # 3. config
    cfg = json.load(open(os.path.join(SRC, "config.json")))
    cfg["quantization_config"] = {
        "quant_method": "nvfp4_aqlm_hybrid",
        "nvfp4": cfg["quantization_config"],
        "aqlm": {"entries": 65536, "group_size": 8},
        "aqlm_layer_books": layer_books,
    }
    json.dump(cfg, open(os.path.join(DST, "config.json"), "w"), indent=2)

    # 4. aux files
    for f in os.listdir(SRC):
        if (
            f.endswith(".json")
            and f not in ("config.json", "model.safetensors.index.json")
            or f.endswith((".txt", ".jinja", ".py", ".md"))
        ):
            shutil.copy2(os.path.join(SRC, f), os.path.join(DST, f))
    print("DONE:", DST)


if __name__ == "__main__":
    main()