#!/usr/bin/env python3 """REAP-style expert saliency scoring, CPU-only (runs alongside GPU jobs). Per expert e of every MoE layer: score_e = SUM_{t routed to e} g_{t,e} * ||f_e(x_t)||_2 * relerr_e g router weight (topk_weights from /data/glm52-acts) f_e(x) down(silu(gate x) * up x) computed with TEACHER (NVFP4-dequant) weights, fp32 on CPU relerr_e mean of w13/w2 h-weighted 2-bit reconstruction rel-errors using the init AQLM parts (/data/glm52-aqlm-parts cover all 256 experts/layer); h = E[x^2] over the expert's routed tokens Teachers: layers 3,4,5,8,74-77 -> /data/glm52-old-layerwise (all experts); other layers: cold experts from /tmp/glm52-hot-dl2 regions, hot experts from the live checkpoint's compacted nvfp4_* arrays. Output: /data/glm52-reap-scores.npz (scores [75,256], raw saliency, relerr, token hits, layer_ids) + per-layer spearman vs frequency counts. Usage: score_experts_reap.py [--layers ...] [--workers 12] """ import argparse import json import os import time from concurrent.futures import ProcessPoolExecutor ACTS = "/data/glm52-acts" CKPT = "/data/glm52" PARTS = "/data/glm52-aqlm-parts" OLD = "/data/glm52-old-layerwise" DL = "/tmp/glm52-hot-dl2" OUTDIR = "/data/glm52-reap-scores" LOCAL_TEACHER_LAYERS = {3, 4, 5, 8, 74, 75, 76, 77} 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 _readers(): import torch from safetensors import safe_open class Regions: def __init__(self): self.headers = json.load(open(f"{DL}/headers.json")) self.wm = {n: s for s, h in self.headers.items() for n in h["header"] if n != "__metadata__"} 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) def get(self, name): DT = {"U8": torch.uint8, "F8_E4M3": torch.uint8, "BF16": torch.bfloat16, "F32": torch.float32} 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=DT[info["dtype"]] ).reshape(info["shape"]) raise KeyError(name) class St: def __init__(self, root): idx = json.load(open(f"{root}/model.safetensors.index.json")) self.wm = idx["weight_map"] self.root = root self._o = {} def get(self, name): shard = self.wm[name] if shard not in self._o: self._o[shard] = safe_open(f"{self.root}/{shard}", framework="pt") return self._o[shard].get_tensor(name) return Regions(), St(OLD), St(CKPT) def _dequant_nvfp4(packed, bscale, scale2, lut): import torch lo = (packed & 0x0F).long() hi = (packed >> 4).long() vals = torch.empty(packed.shape[0], packed.shape[1] * 2, dtype=torch.float32) vals[:, 0::2] = lut[lo] vals[:, 1::2] = lut[hi] scale = bscale.view(torch.float8_e4m3fn).float().repeat_interleave( 16, dim=1) return vals * scale * float(scale2) def process_layer(li, threads): import numpy as np import torch torch.set_num_threads(threads) t0 = time.time() lut = torch.tensor(FP4_LUT, dtype=torch.float32) outp = f"{OUTDIR}/layer_{li}.npz" if os.path.exists(outp): return li, "skip" regions, old, live = _readers() acts = torch.load(f"{ACTS}/acts_layer{li}.pt", map_location="cpu", weights_only=True) x = acts["x"].float() tk = acts["topk_ids"].long() tw = acts["topk_weights"].float() p = f"model.layers.{li}.mlp.experts" kind = live.get(f"{p}.hyb_kind") hot_pos = {int(e): j for j, e in enumerate((kind == 0).nonzero().flatten().tolist())} nv = {n: live.get(f"{p}.{n}") for n in ("nvfp4_w13_packed", "nvfp4_w13_bscale", "nvfp4_w13_scale2", "nvfp4_w2_packed", "nvfp4_w2_bscale", "nvfp4_w2_scale2")} part = torch.load(f"{PARTS}/layer_{li}.pt", map_location="cpu", weights_only=True) cb13 = part["w13_codebooks"][0].float() cb2 = part["w2_codebooks"][0].float() def teacher(e): if li in LOCAL_TEACHER_LAYERS: ep = f"model.layers.{li}.mlp.experts.{e}" g = _dequant_nvfp4(old.get(f"{ep}.gate_proj.weight"), old.get(f"{ep}.gate_proj.weight_scale").view(torch.uint8), old.get(f"{ep}.gate_proj.weight_scale_2"), lut) u = _dequant_nvfp4(old.get(f"{ep}.up_proj.weight"), old.get(f"{ep}.up_proj.weight_scale").view(torch.uint8), old.get(f"{ep}.up_proj.weight_scale_2"), lut) d = _dequant_nvfp4(old.get(f"{ep}.down_proj.weight"), old.get(f"{ep}.down_proj.weight_scale").view(torch.uint8), old.get(f"{ep}.down_proj.weight_scale_2"), lut) return g, u, d if e in hot_pos: j = hot_pos[e] w13 = _dequant_nvfp4(nv["nvfp4_w13_packed"][j], nv["nvfp4_w13_bscale"][j], 1.0, lut) # scale2 is per (gate,up): w13[:2048] *= float(nv["nvfp4_w13_scale2"][j, 0]) w13[2048:] *= float(nv["nvfp4_w13_scale2"][j, 1]) w2 = _dequant_nvfp4(nv["nvfp4_w2_packed"][j], nv["nvfp4_w2_bscale"][j], float(nv["nvfp4_w2_scale2"][j, 0]), lut) return w13[:2048], w13[2048:], w2 ep = f"model.layers.{li}.mlp.experts.{e}" g = _dequant_nvfp4(regions.get(f"{ep}.gate_proj.weight"), regions.get(f"{ep}.gate_proj.weight_scale").view(torch.uint8), regions.get(f"{ep}.gate_proj.weight_scale_2"), lut) u = _dequant_nvfp4(regions.get(f"{ep}.up_proj.weight"), regions.get(f"{ep}.up_proj.weight_scale").view(torch.uint8), regions.get(f"{ep}.up_proj.weight_scale_2"), lut) d = _dequant_nvfp4(regions.get(f"{ep}.down_proj.weight"), regions.get(f"{ep}.down_proj.weight_scale").view(torch.uint8), regions.get(f"{ep}.down_proj.weight_scale_2"), lut) return g, u, d def aqlm_dequant(codes, cb, scales): idx = codes.view(torch.uint16).long() # [books,M,K8] w = cb[idx[0]] return w.reshape(w.shape[0], -1) * scales.unsqueeze(-1) E = 256 sal = np.zeros(E) rel = np.zeros(E) hits = np.zeros(E, dtype=np.int64) counts = np.zeros(E, dtype=np.int64) for e in range(E): mask = (tk == e) rows = mask.any(1) counts[e] = int(rows.sum()) g_w, u_w, d_w = teacher(e) if counts[e] > 0: xe = x[rows] ge = tw[mask] # one weight per hit row (expert unique per token) mid = torch.nn.functional.silu(xe @ g_w.t()) * (xe @ u_w.t()) f = mid @ d_w.t() sal[e] = float((ge * f.norm(dim=1)).sum()) h13 = xe.pow(2).mean(0) h2 = mid.pow(2).mean(0) else: h13 = x.pow(2).mean(0) h2 = None # relerr from init AQLM parts (per-expert, h-weighted) w13_t = torch.cat([g_w, u_w], 0) w13_q = aqlm_dequant(part["w13_codes"][e], cb13, part["w13_scales"][e].float()) num = ((w13_t - w13_q).pow(2) * h13).sum() den = (w13_t.pow(2) * h13).sum().clamp_min(1e-12) r13 = float((num / den).sqrt()) w2_q = aqlm_dequant(part["w2_codes"][e][:1], cb2, part["w2_scales"][e].float()) if h2 is None: h2 = torch.ones(2048) num2 = ((d_w - w2_q).pow(2) * h2).sum() den2 = (d_w.pow(2) * h2).sum().clamp_min(1e-12) r2 = float((num2 / den2).sqrt()) rel[e] = 0.5 * (r13 + r2) hits[e] = counts[e] score = sal * rel from scipy.stats import spearmanr rho = spearmanr(score, counts).statistic if counts.sum() else 0.0 np.savez(outp, score=score, saliency=sal, relerr=rel, hits=hits, counts=counts) return li, f"done {time.time()-t0:.0f}s spearman(score,freq)={rho:.3f}" def main(): ap = argparse.ArgumentParser() ap.add_argument("--layers", default=None) ap.add_argument("--workers", type=int, default=12) ap.add_argument("--threads", type=int, default=9) args = ap.parse_args() os.makedirs(OUTDIR, exist_ok=True) layers = list(range(3, 78)) if args.layers: want = {int(v) for v in args.layers.split(",")} layers = [li for li in layers if li in want] with ProcessPoolExecutor(max_workers=args.workers) as ex: futs = {ex.submit(process_layer, li, args.threads): li for li in layers} for f in futs: pass for f in list(futs): li, msg = f.result() print(f"L{li}: {msg}", flush=True) # merge import numpy as np files = sorted(int(f[6:-4]) for f in os.listdir(OUTDIR) if f.startswith("layer_")) if len(files) == 75: merged = {k: np.stack([np.load(f"{OUTDIR}/layer_{li}.npz")[k] for li in files]) for k in ("score", "saliency", "relerr", "hits", "counts")} np.savez("/data/glm52-reap-scores.npz", layer_ids=np.array(files), **merged) print("merged -> /data/glm52-reap-scores.npz") if __name__ == "__main__": main()