code/tools/aqlm_converge.py

#!/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()