code/tools/phase15_fit.py

#!/usr/bin/env python3
"""Phase 1.5: encode REAP-demoted experts against converged codebooks.

For each layer: experts that are HOT in the live checkpoint but COLD in
any of the three REAP assignments get AQLM codes: h-weighted encode
(2 sweeps) against the layer's CONVERGED codebook (frozen) + weighted
scale refit. Teachers = the live checkpoint's own NVFP4 hot arrays.
Output: /data/glm52-aqlm-conv15/layer_N.pt = conv parts EXTENDED with the
new experts (expert_ids re-sorted ascending).
"""
import json
import os
import time
from concurrent.futures import ProcessPoolExecutor

CKPT = "/data/glm52"
CONV = os.environ.get("P15_CONV", "/data/glm52-aqlm-conv")
OUT = os.environ.get("P15_OUT", "/data/glm52-aqlm-conv15")
ACTS = "/data/glm52-acts"
ASSIGNS = ["/data/glm52-assign-reap-250.json",
           "/data/glm52-assign-reap-290.json",
           "/data/glm52-assign-reap-310.json"]
G = 8
CHUNK = 32768
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]


def process_layer(li, gpu):
    import torch
    from safetensors import safe_open

    torch.set_num_threads(4)
    torch.cuda.set_device(gpu)
    dev = f"cuda:{gpu}"
    t0 = time.time()
    outp = f"{OUT}/layer_{li}.pt"
    if os.path.exists(outp):
        return li, "skip"
    lut = torch.tensor(FP4_LUT, dtype=torch.float32, device=dev)

    idx = json.load(open(f"{CKPT}/model.safetensors.index.json"))
    wm = idx["weight_map"]
    opened = {}

    def ck(name):
        s = wm[f"model.layers.{li}.mlp.experts.{name}"]
        if s not in opened:
            opened[s] = safe_open(f"{CKPT}/{s}", framework="pt")
        return opened[s].get_tensor(f"model.layers.{li}.mlp.experts.{name}")

    kind = ck("hyb_kind")
    cur_hot = (kind == 0).nonzero().flatten().tolist()
    hot_pos = {e: j for j, e in enumerate(cur_hot)}
    need_file = os.environ.get("P15_NEED")
    if need_file:
        need = json.load(open(need_file))
        new = sorted(need.get(str(li), []))
    else:
        union_cold = set()
        for f in ASSIGNS:
            a = json.load(open(f))[str(li)]
            union_cold |= set(a["cold"])
        new = sorted(set(cur_hot) & union_cold)   # demoted: need codes
    part = torch.load(f"{CONV}/layer_{li}.pt", map_location="cpu",
                      weights_only=True)
    if not new:
        torch.save(part, outp)
        return li, "no demotions, copied"

    acts = torch.load(f"{ACTS}/acts_layer{li}.pt", map_location="cpu",
                      weights_only=True)
    x = acts["x"].to(dev).float()
    tk = acts["topk_ids"].to(dev).long()

    nvp = {n: ck(n) for n in ("nvfp4_w13_packed", "nvfp4_w13_bscale",
                              "nvfp4_w13_scale2", "nvfp4_w2_packed",
                              "nvfp4_w2_bscale", "nvfp4_w2_scale2")}

    def teacher(e):
        j = hot_pos[e]
        pk = nvp["nvfp4_w13_packed"][j].to(dev)
        bs = nvp["nvfp4_w13_bscale"][j].to(dev)
        lo = lut[(pk & 0xF).long()]
        hi = lut[(pk >> 4).long()]
        w13 = torch.stack([lo, hi], -1).reshape(pk.shape[0], -1)
        w13 *= bs.view(torch.float8_e4m3fn).float().repeat_interleave(16, -1)
        w13[:2048] *= float(nvp["nvfp4_w13_scale2"][j, 0])
        w13[2048:] *= float(nvp["nvfp4_w13_scale2"][j, 1])
        pk2 = nvp["nvfp4_w2_packed"][j].to(dev)
        bs2 = nvp["nvfp4_w2_bscale"][j].to(dev)
        lo2 = lut[(pk2 & 0xF).long()]
        hi2 = lut[(pk2 >> 4).long()]
        w2 = torch.stack([lo2, hi2], -1).reshape(pk2.shape[0], -1)
        w2 *= bs2.view(torch.float8_e4m3fn).float().repeat_interleave(16, -1)
        w2 *= float(nvp["nvfp4_w2_scale2"][j, 0])
        return w13, w2

    def encode_scaled(w, h, cent):
        """2-sweep h-weighted encode with scale refit; cent frozen."""
        M, K = w.shape
        NG = K // G
        scales = w.abs().mean(-1).clamp_min(1e-8)
        centb = cent.t().to(torch.bfloat16)
        cent2b = (cent.t() ** 2).to(torch.bfloat16)
        hw = h.reshape(1, NG, G).expand(M, NG, G).reshape(-1, G)
        codes = torch.empty(M * NG, dtype=torch.int32, device=dev)
        for _ in range(2):
            tgt = (w / scales.unsqueeze(-1)).reshape(-1, G)
            hwb = hw.to(torch.bfloat16)
            for s in range(0, tgt.shape[0], CHUNK):
                v = tgt[s:s+CHUNK].to(torch.bfloat16)
                ww = hwb[s:s+CHUNK]
                a = ww @ cent2b
                a -= 2 * ((v * ww) @ centb)
                codes[s:s+CHUNK] = a.argmin(-1).to(torch.int32)
            dec = cent[codes.long()].reshape(M, K)
            num = (w * h * dec).sum(-1)
            den = ((dec * dec) * h).sum(-1).clamp_min(1e-10)
            scales = (num / den).clamp(1e-6, None)
        c16 = torch.where(codes >= 32768, codes - 65536, codes).to(torch.int16)
        return c16.reshape(1, M, NG), scales

    cb13 = part["w13_codebooks"][0].float().to(dev)
    cb2 = part["w2c_codebooks"][0].float().to(dev)
    add = {"ids": [], "w13c": [], "w13s": [], "w2c": [], "w2s": []}
    for e in new:
        w13, w2 = teacher(e)
        rows = (tk == e).any(1)
        if int(rows.sum()) >= 16:
            xe = x[rows]
            h13 = xe.pow(2).mean(0).clamp_min(1e-10)
            gate, up = w13[:2048], w13[2048:]
            mid = torch.nn.functional.silu(xe @ gate.t()) * (xe @ up.t())
            h2 = mid.pow(2).mean(0).clamp_min(1e-10)
        else:
            h13 = x.pow(2).mean(0).clamp_min(1e-10)
            h2 = torch.ones(2048, device=dev)
        c13, s13 = encode_scaled(w13, h13, cb13)
        c2, s2 = encode_scaled(w2, h2, cb2)
        add["ids"].append(e)
        add["w13c"].append(c13.unsqueeze(0).cpu())
        add["w13s"].append(s13.half().unsqueeze(0).cpu())
        add["w2c"].append(c2.unsqueeze(0).cpu())
        add["w2s"].append(s2.half().unsqueeze(0).cpu())

    import torch as T
    old_ids = part["expert_ids"].tolist()
    all_ids = old_ids + add["ids"]
    order = sorted(range(len(all_ids)), key=lambda i: all_ids[i])
    cat = {
        "w13_codes": T.cat([part["w13_codes"]] + add["w13c"]),
        "w13_scales": T.cat([part["w13_scales"]] + add["w13s"]),
        "w2c_codes": T.cat([part["w2c_codes"]] + add["w2c"]),
        "w2c_scales": T.cat([part["w2c_scales"]] + add["w2s"]),
    }
    sel = T.tensor(order)
    T.save({
        "layer": li,
        "expert_ids": T.tensor([all_ids[i] for i in order],
                               dtype=T.int32),
        "w13_codes": cat["w13_codes"][sel].contiguous(),
        "w13_codebooks": part["w13_codebooks"],
        "w13_scales": cat["w13_scales"][sel].contiguous(),
        "w2c_codes": cat["w2c_codes"][sel].contiguous(),
        "w2c_codebooks": part["w2c_codebooks"],
        "w2c_scales": cat["w2c_scales"][sel].contiguous(),
        "w13_err_before": part["w13_err_before"],
        "w13_err_after": part["w13_err_after"],
        "w2_err_before": part["w2_err_before"],
        "w2_err_after": part["w2_err_after"],
    }, outp)
    return li, f"+{len(new)} demoted encoded ({time.time()-t0:.0f}s)"


def main():
    os.makedirs(OUT, exist_ok=True)
    layers = range(3, 78)
    if os.environ.get("P15_NEED"):
        layers = sorted(int(k) for k in
                        json.load(open(os.environ["P15_NEED"])))
    jobs = [(li, i % 8) for i, li in enumerate(layers)]
    with ProcessPoolExecutor(max_workers=8) as ex:
        for f in [ex.submit(process_layer, *j) for j in jobs]:
            li, msg = f.result()
            print(f"L{li}: {msg}", flush=True)


if __name__ == "__main__":
    main()