File size: 6,098 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 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 | #!/usr/bin/env python3
"""Assemble the hybrid NVFP4+AQLM GLM-5.2 checkpoint at /data/glm52.
Sources:
- /tmp/glm52-dl/nvfp4-full: lukealonso/GLM-5.2-NVFP4 (all non-expert
tensors verbatim; routed-expert tensors only for plan.nvfp4_layers)
- /data/glm52-aqlm-parts/layer_N.pt: AQLM codes/codebooks/scales for
the remaining expert layers
Output: sharded safetensors (~4GB each) + index + config.json with the
nvfp4_aqlm_hybrid quantization_config + tokenizer/aux files.
"""
import json
import os
import re
import shutil
import time
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"
DST = "/data/glm52"
SHARD_BYTES = 4 << 30
plan = json.load(open(os.path.join(ROOT, "hybrid_plan.json")))
NVFP4_LAYERS = set(plan["nvfp4_layers"])
AQLM_LAYERS = sorted(plan["aqlm_mixed_layers"]) + sorted(plan["aqlm_cold_layers"])
def keep(name: str) -> bool:
m = re.match(r"model\.layers\.(\d+)\.mlp\.experts\.\d+\.", 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()
# rename to canonical model-XXXXX-of-NNNNN scheme
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
idx = {
"metadata": {"total_size": self.total},
"weight_map": wm,
}
json.dump(
idx,
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")
def main():
os.makedirs(DST, exist_ok=True)
idx = json.load(open(os.path.join(SRC, "model.safetensors.index.json")))
weight_map = idx["weight_map"]
shards = sorted(set(weight_map.values()))
writer = ShardWriter(DST)
# 1. copy kept tensors from the NVFP4 repo, shard by shard
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
print(f"{shard}: keeping {len(names)} tensors", flush=True)
for n in names:
writer.add(n, f.get_tensor(n))
# 2. append AQLM expert tensors
layer_books = {}
mse_report = {}
for li in AQLM_LAYERS:
part = os.path.join(PARTS, f"layer_{li}.pt")
# the quantizer writes in place: wait for existence + 60s of quiescence
while (
not os.path.exists(part) or time.time() - os.path.getmtime(part) < 60
):
print(f"waiting for quantizer: layer {li} ...", flush=True)
time.sleep(30)
d = torch.load(part, map_location="cpu", weights_only=True)
p = f"model.layers.{li}.mlp.experts"
writer.add(f"{p}.w13_codes", d["w13_codes"])
writer.add(f"{p}.w13_codebooks", d["w13_codebooks"])
writer.add(f"{p}.w13_scales", d["w13_scales"])
writer.add(f"{p}.w2_codes", d["w2_codes"])
writer.add(f"{p}.w2_codebooks", d["w2_codebooks"])
writer.add(f"{p}.w2_scales", d["w2_scales"])
layer_books[str(li)] = {
"w13": d["books"]["w13"],
"w2": d["books"]["w2"],
}
mse_report[li] = (d["w13_rel_mse"], d["w2_rel_mse"])
print(
f"layer {li}: aqlm books={d['books']} "
f"rel_mse w13={d['w13_rel_mse']:.4f} w2={d['w2_rel_mse']:.4f}",
flush=True,
)
writer.finalize()
# 3. config.json: replace quantization_config with the hybrid config,
# dropping quantized expert layers from the modelopt view is NOT
# needed -- AQLM layers are dispatched before the nvfp4 config is
# consulted, and their tensors do not match NVFP4 patterns.
cfg = json.load(open(os.path.join(SRC, "config.json")))
nvfp4_qc = cfg["quantization_config"]
cfg["quantization_config"] = {
"quant_method": "nvfp4_aqlm_hybrid",
"nvfp4": nvfp4_qc,
"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(f"copied {f}")
json.dump(
{str(k): v for k, v in mse_report.items()},
open(os.path.join(DST, "aqlm_mse_report.json"), "w"),
indent=2,
)
print("DONE:", DST)
if __name__ == "__main__":
main()
|