File size: 2,584 Bytes
fdc6474 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 | #!/usr/bin/env python3
"""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))
# drop old per-expert NVFP4 entries for this layer
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}")
# recompute logical total size
json.dump(idx, open(idx_path, "w"), indent=0)
json.dump(cfg, open(os.path.join(DST, "config.json"), "w"), indent=2)
# summarize logical size (sum over indexed tensors is expensive; report delta)
print("index now has", len(wm), "tensors")
print("aqlm_layer_books entries:", len(books_cfg))
print("DONE")
|