code/tools/build_checkpoint_v6.py

#!/usr/bin/env python3
"""Build a two-tier variant checkpoint with per-expert source dispatch.

Hot NVFP4 bytes come from either the v4-uniform checkpoint's compacted
arrays (experts already hot there) or the ranged downloads in
/tmp/glm52-hot-dl2 (newly-hot experts). Cold AQLM slices come from
/data/glm52-aqlm-parts. Non-expert tensors stream from v4-uniform.

Usage: build_checkpoint_v6.py <assignment.json> <dst_dir>
"""

import json
import os
import re
import shutil
import sys

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

SRC = "/data/glm52-v4-uniform"
DL = "/tmp/glm52-hot-dl2"
PARTS = "/data/glm52-aqlm-parts"
ASSIGN = sys.argv[1]
DST = sys.argv[2]
SHARD_BYTES = 4 << 30
N_EXP, INTER, HIDDEN = 256, 2048, 6144

assignment = {int(k): v for k, v in json.load(open(ASSIGN)).items()}

DTYPES = {"U8": torch.uint8, "F8_E4M3": torch.uint8, "BF16": torch.bfloat16,
          "F32": torch.float32, "F16": torch.float16, "I16": torch.int16,
          "I8": torch.int8}


class RegionReader:
    def __init__(self):
        self.headers = json.load(open(f"{DL}/headers.json"))
        self.wm = {}
        for shard, h in self.headers.items():
            for name in h["header"]:
                if name != "__metadata__":
                    self.wm[name] = shard
        self.regions = {}
        for shard in self.headers:
            d = f"{DL}/regions/{shard}"
            regs = []
            if os.path.isdir(d):
                for f in os.listdir(d):
                    p = os.path.join(d, f)
                    regs.append((int(f[:-4]), p, os.path.getsize(p)))
            self.regions[shard] = sorted(regs)

    def get(self, name):
        shard = self.wm[name]
        info = self.headers[shard]["header"][name]
        b, e = info["data_offsets"]
        for rb, path, sz in self.regions[shard]:
            if rb <= b and e <= rb + sz:
                with open(path, "rb") as fh:
                    fh.seek(b - rb)
                    buf = fh.read(e - b)
                t = torch.frombuffer(bytearray(buf), dtype=DTYPES[info["dtype"]])
                return t.reshape(info["shape"])
        raise KeyError(f"{name}: not in downloaded regions")


class SrcReader:
    def __init__(self):
        idx = json.load(open(f"{SRC}/model.safetensors.index.json"))
        self.wm = idx["weight_map"]
        self._open = {}

    def get(self, name):
        shard = self.wm[name]
        if shard not in self._open:
            self._open[shard] = safe_open(f"{SRC}/{shard}", framework="pt")
        return self._open[shard].get_tensor(name)


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

    def add(self, name, 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
        self.files.append(fname)
        self.cur, self.cur_bytes = {}, 0

    def finalize(self):
        self.flush()
        wm = {}
        for i, fname in enumerate(self.files, 1):
            new = f"model-{i:05d}-of-{self.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(f"{self.dst}/model.safetensors.index.json", "w"), indent=0)
        print(f"index: {len(wm)} tensors, {self.total/1e9:.1f} GB", flush=True)


def main():
    os.makedirs(DST, exist_ok=True)
    sr = SrcReader()
    rr = RegionReader()
    writer = ShardWriter(DST)

    exp_pat = re.compile(r"model\.layers\.(\d+)\.mlp\.experts\.")
    for shard in sorted(set(sr.wm.values())):
        with safe_open(f"{SRC}/{shard}", framework="pt") as f:
            for n in f.keys():
                m = exp_pat.match(n)
                if m and int(m.group(1)) != 78:  # keep MTP verbatim
                    continue
                writer.add(n, f.get_tensor(n))

    layer_books = {}
    for li in sorted(assignment):
        hot = sorted(assignment[li]["hot"])
        cold = sorted(assignment[li]["cold"])
        assert len(hot) + len(cold) == N_EXP
        p = f"model.layers.{li}.mlp.experts"

        kind_old = sr.get(f"{p}.hyb_kind")
        old_hot = (kind_old == 0).nonzero().flatten().tolist()
        pos = {e: j for j, e in enumerate(old_hot)}
        old = {n: sr.get(f"{p}.{n}") for n in
               ("nvfp4_w13_packed", "nvfp4_w13_bscale", "nvfp4_w13_scale2",
                "nvfp4_w2_packed", "nvfp4_w2_bscale", "nvfp4_w2_scale2")}

        na = len(hot)
        w13p = torch.empty(na, 2*INTER, HIDDEN//2, dtype=torch.uint8)
        w13b = torch.empty(na, 2*INTER, HIDDEN//16, dtype=torch.uint8)
        w13s = torch.empty(na, 2, dtype=torch.float32)
        w2p = torch.empty(na, HIDDEN, INTER//2, dtype=torch.uint8)
        w2b = torch.empty(na, HIDDEN, INTER//16, dtype=torch.uint8)
        w2s = torch.empty(na, 1, dtype=torch.float32)
        n_dl = 0
        for j, e in enumerate(hot):
            if e in pos:
                k = pos[e]
                w13p[j] = old["nvfp4_w13_packed"][k]
                w13b[j] = old["nvfp4_w13_bscale"][k]
                w13s[j] = old["nvfp4_w13_scale2"][k]
                w2p[j] = old["nvfp4_w2_packed"][k]
                w2b[j] = old["nvfp4_w2_bscale"][k]
                w2s[j] = old["nvfp4_w2_scale2"][k]
            else:
                n_dl += 1
                ep = f"model.layers.{li}.mlp.experts.{e}"
                w13p[j, :INTER] = rr.get(f"{ep}.gate_proj.weight")
                w13p[j, INTER:] = rr.get(f"{ep}.up_proj.weight")
                w2p[j] = rr.get(f"{ep}.down_proj.weight")
                w13b[j, :INTER] = rr.get(
                    f"{ep}.gate_proj.weight_scale").view(torch.uint8)
                w13b[j, INTER:] = rr.get(
                    f"{ep}.up_proj.weight_scale").view(torch.uint8)
                w2b[j] = rr.get(
                    f"{ep}.down_proj.weight_scale").view(torch.uint8)
                w13s[j, 0] = rr.get(f"{ep}.gate_proj.weight_scale_2").float()
                w13s[j, 1] = rr.get(f"{ep}.up_proj.weight_scale_2").float()
                w2s[j, 0] = rr.get(f"{ep}.down_proj.weight_scale_2").float()

        kind = torch.full((N_EXP,), 2, dtype=torch.int8)
        for e in hot:
            kind[e] = 0
        part = torch.load(f"{PARTS}/layer_{li}.pt", map_location="cpu",
                          weights_only=True)
        cold_idx = torch.tensor(cold, dtype=torch.long)
        writer.add(f"{p}.hyb_kind", kind)
        writer.add(f"{p}.w13_codes", part["w13_codes"][cold_idx].contiguous())
        writer.add(f"{p}.w13_codebooks", part["w13_codebooks"].clone())
        writer.add(f"{p}.w13_scales", part["w13_scales"][cold_idx].contiguous())
        writer.add(f"{p}.w2m_codes",
                   torch.empty(0, 2, HIDDEN, INTER//8, dtype=torch.int16))
        writer.add(f"{p}.w2m_codebooks", part["w2_codebooks"].clone())
        writer.add(f"{p}.w2m_scales", torch.empty(0, HIDDEN, dtype=torch.float16))
        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())
        for n, t in zip(("nvfp4_w13_packed", "nvfp4_w13_bscale",
                         "nvfp4_w13_scale2", "nvfp4_w2_packed",
                         "nvfp4_w2_bscale", "nvfp4_w2_scale2"),
                        (w13p, w13b, w13s, w2p, w2b, w2s)):
            writer.add(f"{p}.{n}", t)
        layer_books[str(li)] = {"n_nvfp4": na, "n_base": 0, "n_cold": len(cold)}
        print(f"layer {li}: hot={na} ({n_dl} from download) cold={len(cold)}",
              flush=True)

    writer.finalize()
    cfg = json.load(open(f"{SRC}/config.json"))
    cfg["quantization_config"]["aqlm_layer_books"] = layer_books
    json.dump(cfg, open(f"{DST}/config.json", "w"), indent=2)
    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(f"{SRC}/{f}", f"{DST}/{f}")
    print("DONE:", DST)


if __name__ == "__main__":
    main()