code/tools/build_checkpoint_v5.py

#!/usr/bin/env python3
"""Build the two-tier (hot NVFP4 / cold 2-bpw AQLM) checkpoint at /data/glm52-v4.

Sources:
  - /data/glm52 (v3): all non-expert tensors; compacted NVFP4 hot arrays for
    LOCAL_LAYERS (their new hot sets are subsets of the stored ones); the
    MTP layer's per-expert tensors (copied verbatim).
  - /tmp/glm52-hot-dl: ranged NVFP4 regions for the other layers' hot experts.
  - /data/glm52-aqlm-parts/layer_N.pt: w13 1-book codes (all experts) and
    2-book w2 codes; cold w2 takes book 0.
"""

import json
import os
import re
import shutil

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

SRC = "/data/glm52"
DL = "/tmp/glm52-hot-dl"
PARTS = "/data/glm52-aqlm-parts"
ASSIGN = "/data/glm52-expert-assignment.json"
DST = "/data/glm52-v5"
SHARD_BYTES = 4 << 30
LOCAL_LAYERS = set(range(3, 78))
N_EXP, INTER, HIDDEN = 256, 2048, 6144

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

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:
    """Read tensors from the ranged-download regions."""

    def __init__(self):
        self.headers = json.load(open(f"{DL}/headers.json"))
        idx = json.load(open(f"{DL}/index.json")) if os.path.exists(
            f"{DL}/index.json") else None
        self.wm = {}
        for shard, h in self.headers.items():
            for name in h["header"]:
                if name != "__metadata__":
                    self.wm[name] = shard
        # region files per shard: sorted list of (rel_start, path, size)
        self.regions = {}
        for shard in self.headers:
            d = f"{DL}/regions/{shard}"
            regs = []
            if os.path.isdir(d):
                for f in os.listdir(d):
                    regs.append((int(f[:-4]), os.path.join(d, f),
                                 os.path.getsize(os.path.join(d, f))))
            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}: bytes [{b},{e}) 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)
        print(f"  wrote {fname} ({self.cur_bytes/1e9:.2f} GB)", flush=True)
        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")


def hot_arrays_from_download(li, hot, rr):
    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)
    for j, e in enumerate(hot):
        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()
    return w13p, w13b, w13s, w2p, w2b, w2s


def hot_arrays_from_local(li, hot, sr, cfg_books):
    """Slice the stored compacted arrays (old hot superset, asc expert id)."""
    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)}
    sel = torch.tensor([pos[e] for e in hot], dtype=torch.long)
    return tuple(
        sr.get(f"{p}.{n}")[sel].contiguous()
        for n in ("nvfp4_w13_packed", "nvfp4_w13_bscale", "nvfp4_w13_scale2",
                  "nvfp4_w2_packed", "nvfp4_w2_bscale", "nvfp4_w2_scale2")
    )


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

    # 1. non-expert tensors + MTP per-expert tensors, streamed
    drop = re.compile(
        r"model\.layers\.(\d+)\.mlp\.experts\.(?!78)")  # placeholder, fixed below
    keep_expert_layer = {78}
    shards = sorted(set(sr.wm.values()))
    exp_pat = re.compile(r"model\.layers\.(\d+)\.mlp\.experts\.")
    for shard in shards:
        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)) not in keep_expert_layer:
                    continue
                writer.add(n, f.get_tensor(n))
        print(f"{shard}: copied non-expert tensors", flush=True)

    # 2. hybrid layers, two-tier
    layer_books = {}
    for li in HYBRID:
        hot = sorted(assignment[li]["hot"])
        cold = sorted(assignment[li]["cold"])
        assert len(hot) + len(cold) == N_EXP
        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)
        p = f"model.layers.{li}.mlp.experts"
        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())

        if li in LOCAL_LAYERS:
            arrays = hot_arrays_from_local(li, hot, sr, None)
        else:
            arrays = hot_arrays_from_download(li, hot, rr)
        for n, t in zip(("nvfp4_w13_packed", "nvfp4_w13_bscale",
                         "nvfp4_w13_scale2", "nvfp4_w2_packed",
                         "nvfp4_w2_bscale", "nvfp4_w2_scale2"), arrays):
            writer.add(f"{p}.{n}", t)
        layer_books[str(li)] = {"n_nvfp4": len(hot), "n_base": 0,
                                "n_cold": len(cold)}
        print(f"layer {li}: hot={len(hot)} cold={len(cold)} "
              f"({'local' if li in LOCAL_LAYERS else 'download'})", 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()