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