| |
| """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"]) |
| 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) |
|
|
| 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()) |
| |
| 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()) |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
| |
| layer_books = {} |
| for li in HYBRID_LAYERS: |
| layer_books[str(li)] = build_hybrid_layer(li, reader, writer) |
|
|
| writer.finalize() |
|
|
| |
| 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) |
|
|
| |
| 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() |
|
|
