#!/usr/bin/env python3 """Converged activation-aware AQLM for the cold experts of GLM-5.2. Per (layer, projection), over the layer's cold experts: objective: sum_e || (W_e - What_e) diag(sqrt(h_e)) ||_F^2 with h_e = per-expert diagonal input Hessian E[x^2] from real routed calibration activations (w13: layer input x; w2: the expert's own silu(gate)*up computed with TEACHER weights) variables: shared 65536x8 fp16 codebook (1 book, groups of 8 along the input dim), int16 codes per group, per-expert per-out-channel fp16 scales algorithm: warm start from the existing shipped codes; iterate [weighted coordinate-descent re-encode (batched GEMM argmax under the h-weighted metric) -> closed-form weighted codebook update -> closed-form weighted scale refit] until the relative weighted error improvement < 0.2% or MAX_ITERS. Teacher weights: original NVFP4 per-expert tensors, dequantized to fp32 (from /tmp/glm52-hot-dl2 ranged downloads; layers 3,4,5,8,74-77 from /data/glm52-old-layerwise which stores all 256 experts). Output: /data/glm52-aqlm-conv/layer_N.pt with expert_ids int32 [nC] (ascending), w13_codes int16 [nC,1,4096,768], w13_codebooks fp16 [1,65536,8], w13_scales fp16 [nC,4096], w2c_* likewise ([nC,1,6144,256] codes), plus before/after weighted rel-err metrics. Usage: aqlm_converge.py [--layers a,b,...] [--gpus 0,..7] """ import argparse import json import os import time from concurrent.futures import ProcessPoolExecutor ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) CKPT = "/data/glm52" OLD_LAYERWISE = "/data/glm52-old-layerwise" DL = "/tmp/glm52-hot-dl2" ACTS = "/data/glm52-acts" OUT = "/data/glm52-aqlm-conv" LOCAL_TEACHER_LAYERS = {3, 4, 5, 8, 74, 75, 76, 77} ENTRIES = 65536 G = 8 MAX_ITERS = 6 TOL = 2e-3 CHUNK = 16384 # groups per scoring GEMM chunk FP4_LUT = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0] # ---------------------------------------------------------------- teachers class TeacherReader: """Original NVFP4 per-expert tensors from download regions or the old layer-wise checkpoint.""" DT = None def __init__(self): import torch TeacherReader.DT = {"U8": torch.uint8, "F8_E4M3": torch.uint8, "BF16": torch.bfloat16, "F32": torch.float32} self.headers = json.load(open(f"{DL}/headers.json")) self.wm = {} for shard, h in self.headers.items(): for name in h["header"]: if name != "__metadata__": self.wm[name] = shard self.regions = {} for shard in self.headers: d = f"{DL}/regions/{shard}" regs = [] if os.path.isdir(d): for f in os.listdir(d): p = os.path.join(d, f) regs.append((int(f[:-4]), p, os.path.getsize(p))) self.regions[shard] = sorted(regs) idx = json.load( open(f"{OLD_LAYERWISE}/model.safetensors.index.json")) self.old_wm = idx["weight_map"] self._old_open = {} def _from_regions(self, name): import torch shard = self.wm[name] info = self.headers[shard]["header"][name] b, e = info["data_offsets"] for rb, path, sz in self.regions[shard]: if rb <= b and e <= rb + sz: with open(path, "rb") as fh: fh.seek(b - rb) buf = fh.read(e - b) return torch.frombuffer( bytearray(buf), dtype=self.DT[info["dtype"]] ).reshape(info["shape"]) raise KeyError(name) def _from_old(self, name): from safetensors import safe_open shard = self.old_wm[name] if shard not in self._old_open: self._old_open[shard] = safe_open( f"{OLD_LAYERWISE}/{shard}", framework="pt") return self._old_open[shard].get_tensor(name) def get(self, li, e, proj, part): name = f"model.layers.{li}.mlp.experts.{e}.{proj}.{part}" if li in LOCAL_TEACHER_LAYERS: return self._from_old(name) return self._from_regions(name) def dequant_teacher(tr, li, e, proj, device, lut): import torch packed = tr.get(li, e, proj, "weight").to(device) ws = tr.get(li, e, proj, "weight_scale").view(torch.uint8).to(device) ws2 = tr.get(li, e, proj, "weight_scale_2").to(device).float() lo = (packed & 0x0F).long() hi = (packed >> 4).long() vals = torch.empty(packed.shape[0], packed.shape[1] * 2, dtype=torch.float32, device=device) vals[:, 0::2] = lut[lo] vals[:, 1::2] = lut[hi] scale = ws.view(torch.float8_e4m3fn).float().repeat_interleave(16, dim=1) return vals * scale * ws2 # ------------------------------------------------------------- optimization def weighted_encode(vecs, hw, cent, out_codes): """CD re-encode: for each group i pick argmin_c sum_j hw[i,j]*(v-c)^2. vecs,hw: [N,8] fp32 (hw = per-group diag weights). cent: [K,8] fp32.""" import torch for s in range(0, vecs.shape[0], CHUNK): v = vecs[s:s + CHUNK] w = hw[s:s + CHUNK] # cost = sum_j w_j c_j^2 - 2 sum_j w_j v_j c_j (+const) a = w @ (cent.t() ** 2) # [n,K] b = (v * w) @ cent.t() # [n,K] out_codes[s:s + CHUNK] = (a - 2 * b).argmin(-1).to(torch.int32) return out_codes def codebook_update(vecs, hw, codes, K): """Closed-form weighted LS per entry per dim: c_kj = sum_{i in k} w_ij v_ij / sum w_ij.""" import torch num = torch.zeros(K, G, device=vecs.device) den = torch.zeros(K, G, device=vecs.device) idx = codes.long() num.index_add_(0, idx, vecs * hw) den.index_add_(0, idx, hw) cent = num / den.clamp_min(1e-12) dead = den.sum(-1) < 1e-12 return cent, dead def fit_projection(w, h, codes0, cb0, scales0, gen, log, tag): """w [E,M,K] fp16 teacher; h [E,K] diag Hessian; warm-start codes/cb/scales. Expert-chunked to bound peak memory. Returns codes int16 [E,1,M,K/8], cb fp16 [1,65536,8], scales fp16 [E,M].""" import torch E, M, K = w.shape NG = K // G dev = w.device EC = max(1, min(48, E)) # experts per pass scales = scales0.float().clamp_min(1e-8) cent = cb0[0].float().clone() codes = (codes0.reshape(E, M, NG).to(torch.int32) & 0xFFFF).clone() def chunk_tgt(e0, e1): tgt = (w[e0:e1].float() / scales[e0:e1].unsqueeze(-1)) hh = (h[e0:e1].reshape(e1 - e0, 1, NG, G) .expand(e1 - e0, M, NG, G)) return tgt.reshape(e1 - e0, M, NG, G), hh def weighted_err(): num = torch.zeros((), device=dev, dtype=torch.float64) den = torch.zeros((), device=dev, dtype=torch.float64) for e0 in range(0, E, 16): e1 = min(e0 + 16, E) tgt, hh = chunk_tgt(e0, e1) dec = cent[codes[e0:e1].long().reshape(-1)].reshape( e1 - e0, M, NG, G) num += ((tgt - dec).pow(2) * hh).sum().double() den += (tgt.pow(2) * hh).sum().double() del tgt, hh, dec return (num / den.clamp_min(1e-12)).item() e0_err = weighted_err() prev = e0_err for it in range(MAX_ITERS): num_cb = torch.zeros(ENTRIES, G, device=dev) den_cb = torch.zeros(ENTRIES, G, device=dev) for e0 in range(0, E, EC): e1 = min(e0 + EC, E) tgt, hh = chunk_tgt(e0, e1) v = tgt.reshape(-1, G) hwv = hh.reshape(-1, G) cflat = codes[e0:e1].reshape(-1).clone() weighted_encode(v, hwv, cent, cflat) codes[e0:e1] = cflat.reshape(e1 - e0, M, NG) idx = cflat.long() num_cb.index_add_(0, idx, v * hwv) den_cb.index_add_(0, idx, hwv) del tgt, hh, v, hwv, cflat torch.cuda.empty_cache() newc = num_cb / den_cb.clamp_min(1e-12) dead = den_cb.sum(-1) < 1e-12 newc[dead] = cent[dead] cent = newc # weighted per-channel scale refit, chunked for e0 in range(0, E, EC): e1 = min(e0 + EC, E) dec = cent[codes[e0:e1].long().reshape(-1)].reshape( e1 - e0, M, K) wf = w[e0:e1].float() hh = h[e0:e1].unsqueeze(1) num = (wf * hh * dec).sum(-1) den = ((dec * dec) * hh).sum(-1).clamp_min(1e-10) scales[e0:e1] = (num / den).clamp(1e-6, None) del dec, wf, hh torch.cuda.empty_cache() cur = weighted_err() log(f" {tag} iter{it}: weighted rel-err {cur:.5f}") if prev - cur < TOL * prev: break prev = cur codes16 = torch.where(codes >= 32768, codes - 65536, codes).to(torch.int16) return (codes16.reshape(E, 1, M, NG).cpu(), cent.half().unsqueeze(0).cpu(), scales.half().cpu(), e0_err, prev) def process_layer(li, gpu): import torch t0 = time.time() dev = f"cuda:{gpu}" torch.cuda.set_device(gpu) gen = torch.Generator(device=dev) gen.manual_seed(777 + li) lut = torch.tensor(FP4_LUT, dtype=torch.float32, device=dev) def log(m): print(f"[L{li} gpu{gpu}] {m}", flush=True) outp = os.path.join(OUT, f"layer_{li}.pt") if os.path.exists(outp): log("done, skip") return li, None # cold set + shipped warm start from the live checkpoint from safetensors import safe_open idx = json.load(open(f"{CKPT}/model.safetensors.index.json")) wm = idx["weight_map"] p = f"model.layers.{li}.mlp.experts" opened = {} def ck(name): shard = wm[f"{p}.{name}"] if shard not in opened: opened[shard] = safe_open(f"{CKPT}/{shard}", framework="pt") return opened[shard].get_tensor(f"{p}.{name}") kind = ck("hyb_kind") cold = (kind == 2).nonzero().flatten().tolist() nC = len(cold) # activations -> per-expert diag Hessians acts = torch.load(f"{ACTS}/acts_layer{li}.pt", map_location="cpu", weights_only=True) x = acts["x"].to(dev).float() # [T,6144] tk_ids = acts["topk_ids"].to(dev).long() # [T,8] h13 = torch.zeros(nC, 6144, device=dev) h2 = torch.zeros(nC, 2048, device=dev) pos = {e: j for j, e in enumerate(cold)} tr = TeacherReader() log(f"cold={nC}; dequant teachers + Hessians") w13 = torch.empty(nC, 4096, 6144, dtype=torch.float16, device=dev) w2 = torch.empty(nC, 6144, 2048, dtype=torch.float16, device=dev) hits = (tk_ids.unsqueeze(-1) == torch.tensor(cold, device=dev).view(1, 1, -1)) # [T,8,nC] tok_of = hits.any(1) # [T,nC] for j, e in enumerate(cold): gate = dequant_teacher(tr, li, e, "gate_proj", dev, lut) up = dequant_teacher(tr, li, e, "up_proj", dev, lut) down = dequant_teacher(tr, li, e, "down_proj", dev, lut) w13[j, :2048] = gate.half() w13[j, 2048:] = up.half() w2[j] = down.half() xe = x[tok_of[:, j]] if xe.shape[0] < 16: # rarely-routed: uniform H h13[j] = x.pow(2).mean(0) mid = torch.nn.functional.silu(x[:512] @ gate.t()) * (x[:512] @ up.t()) else: h13[j] = xe.pow(2).mean(0) mid = torch.nn.functional.silu(xe @ gate.t()) * (xe @ up.t()) h2[j] = mid.float().pow(2).mean(0).clamp_min(1e-10) h13[j] = h13[j].clamp_min(1e-10) del hits, tok_of, x torch.cuda.empty_cache() # warm starts (shipped arrays are cold-only, ascending == our order) c13 = ck("w13_codes").to(dev).view(torch.int16) cb13 = ck("w13_codebooks").to(dev) s13 = ck("w13_scales").to(dev) c2 = ck("w2c_codes").to(dev) cb2 = ck("w2c_codebooks").to(dev) s2 = ck("w2c_scales").to(dev) log("fit w13") r13 = fit_projection(w13, h13, c13.view(torch.uint16).int(), cb13.float(), s13, gen, log, "w13") del w13 torch.cuda.empty_cache() log("fit w2") r2 = fit_projection(w2, h2, c2.view(torch.uint16).int(), cb2.float(), s2, gen, log, "w2") del w2 torch.cuda.empty_cache() torch.save({ "layer": li, "expert_ids": torch.tensor(cold, dtype=torch.int32), "w13_codes": r13[0], "w13_codebooks": r13[1], "w13_scales": r13[2], "w2c_codes": r2[0], "w2c_codebooks": r2[1], "w2c_scales": r2[2], "w13_err_before": r13[3], "w13_err_after": r13[4], "w2_err_before": r2[3], "w2_err_after": r2[4], }, outp) log(f"saved ({time.time()-t0:.0f}s) w13 {r13[3]:.4f}->{r13[4]:.4f} " f"w2 {r2[3]:.4f}->{r2[4]:.4f}") return li, (r13[3], r13[4], r2[3], r2[4]) def main(): ap = argparse.ArgumentParser() ap.add_argument("--layers", default=None) ap.add_argument("--gpus", default="0,1,2,3,4,5,6,7") args = ap.parse_args() os.makedirs(OUT, exist_ok=True) layers = list(range(3, 78)) if args.layers: want = {int(x) for x in args.layers.split(",")} layers = [li for li in layers if li in want] gpus = [int(g) for g in args.gpus.split(",")] jobs = [(li, gpus[i % len(gpus)]) for i, li in enumerate(layers)] with ProcessPoolExecutor(max_workers=len(gpus)) as ex: for f in [ex.submit(process_layer, *j) for j in jobs]: li, m = f.result() if m: print(f"== L{li}: w13 {m[0]:.4f}->{m[1]:.4f} " f"w2 {m[2]:.4f}->{m[3]:.4f}", flush=True) if __name__ == "__main__": main()