#!/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-v4" SHARD_BYTES = 4 << 30 LOCAL_LAYERS = {3, 4, 5, 8, 74, 75, 76, 77} 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 = RegionReader() 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()