#!/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 """ 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()