| |
| """Patch /data/glm52 in place after demoting layers from NVFP4 to AQLM. |
| |
| For each demoted layer: append its AQLM tensors as new shards, drop the |
| old per-expert NVFP4 tensor entries from the index (dead bytes stay in the |
| old shards, which is harmless), and update config.json aqlm_layer_books. |
| """ |
|
|
| import json |
| import os |
|
|
| import torch |
| from safetensors.torch import save_file |
|
|
| ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) |
| PARTS = "/data/glm52-aqlm-parts" |
| DST = "/data/glm52" |
|
|
| plan = json.load(open(os.path.join(ROOT, "hybrid_plan.json"))) |
| cfg = json.load(open(os.path.join(DST, "config.json"))) |
| books_cfg = cfg["quantization_config"]["aqlm_layer_books"] |
|
|
| aqlm_all = set(plan["aqlm_mixed_layers"]) | set(plan["aqlm_cold_layers"]) |
| demoted = sorted(li for li in aqlm_all if str(li) not in books_cfg) |
| print("layers to patch:", demoted) |
| assert demoted, "nothing to patch" |
|
|
| idx_path = os.path.join(DST, "model.safetensors.index.json") |
| idx = json.load(open(idx_path)) |
| wm = idx["weight_map"] |
|
|
| for li in demoted: |
| part = torch.load( |
| os.path.join(PARTS, f"layer_{li}.pt"), map_location="cpu", |
| weights_only=True, |
| ) |
| prefix = f"model.layers.{li}.mlp.experts" |
| tensors = { |
| f"{prefix}.w13_codes": part["w13_codes"], |
| f"{prefix}.w13_codebooks": part["w13_codebooks"], |
| f"{prefix}.w13_scales": part["w13_scales"], |
| f"{prefix}.w2_codes": part["w2_codes"], |
| f"{prefix}.w2_codebooks": part["w2_codebooks"], |
| f"{prefix}.w2_scales": part["w2_scales"], |
| } |
| fname = f"model-aqlm-patch-layer{li}.safetensors" |
| save_file(tensors, os.path.join(DST, fname)) |
|
|
| |
| stale = [n for n in wm if n.startswith(f"{prefix}.") and ".w13_" not in n |
| and ".w2_" not in n] |
| for n in stale: |
| del wm[n] |
| for n in tensors: |
| wm[n] = fname |
| books_cfg[str(li)] = { |
| "w13": part["books"]["w13"], "w2": part["books"]["w2"], |
| } |
| nb = sum(t.numel() * t.element_size() for t in tensors.values()) |
| print(f"layer {li}: -{len(stale)} nvfp4 tensors, +{fname} " |
| f"({nb/1e9:.2f} GB), rel_mse w13={part['w13_rel_mse']:.4f} " |
| f"w2={part['w2_rel_mse']:.4f}") |
|
|
| |
| json.dump(idx, open(idx_path, "w"), indent=0) |
| json.dump(cfg, open(os.path.join(DST, "config.json"), "w"), indent=2) |
|
|
| |
| print("index now has", len(wm), "tensors") |
| print("aqlm_layer_books entries:", len(books_cfg)) |
| print("DONE") |
|
|
