#!/usr/bin/env python3 """Codec-v2 encode of one GLM-5.3-Flash routed layer (or an expert range of it). Same runtime format as r27. Changes versus the r27 encoder (glm53_nvfp4/quantize_p8_coupled_rate_layer.py), and nothing else: 1. Hessian: EVERY routed token of the expert in the 64 fit windows (route^2-weighted, exactly as output_aware.route_weighted_hessian), mixed trace-normalised 50/50 (beta) with an all-token Hessian (every 4th token of the 64 fit windows pushed through the expert, unweighted). r27: first 256 routed tokens. 2. GPTQ damping: percdamp = 0.3 x mean diagonal. r27: 0.01. 3. UE8M0 block scales: nearest power of two of the continuous least-squares block scale found with two trellis passes, then the unchanged r27 loop (3 passes, 2 power-of-two refits). r27: scalar nearest-level rule. Transform (coupled H512/H128, sign draw 0), EXL3 suh/svh (read from the r27 sidecars: byte-identical), the procedural MCG alpha-2 E4M3 law, trellis packing and UE8M0 packing are exactly r27's; the r27 functions are imported read-only. Closure check per projection: the packed payload must decode to the encoder reconstruction. """ from __future__ import annotations import argparse import hashlib import json import os import time from pathlib import Path import numpy as np import torch from safetensors import safe_open from safetensors.torch import save_file from glm53_nvfp4 import trellis_mxf as tm from glm53_nvfp4 import trellis_nvfp4 as tn from glm53_nvfp4.capture import LayerCapture from glm53_nvfp4.p8_coupled_scale import ( COUPLED_SIGN_DRAW, CoupledScaleSet, coupled_input_carrier, coupled_middle_carrier, encode_coupled_scale_weights, ) from glm53_nvfp4.shard_index import IndexedCheckpoint CAPTURE = Path(os.environ["CV2_CAPTURE"]) ROLES = Path(os.environ["CV2_ROLES"]) SOURCE = Path(os.environ["CV2_BF16"]) R27 = Path(os.environ["CV2_R27"]) PROJ = ("gate_proj", "up_proj", "down_proj") def r27_scales(layer: int) -> CoupledScaleSet: """EXL3 suh/svh exactly as stored in the r27 rank sidecars (verified byte-identical to the EXL3 checkpoint).""" suh = svh = None parts = [] for r in range(4): with safe_open(str(R27 / "sidecars" / f"p8-layer-{layer:03d}-tp4-rank-{r}.safetensors"), framework="pt", device="cpu") as s: a, b, inter = s.get_tensor("gate_up_suh_fp16"), s.get_tensor("down_svh_fp16"), s.get_tensor( "intermediate_scales_fp16") if suh is None: suh, svh = a, b elif not (torch.equal(suh, a) and torch.equal(svh, b)): raise RuntimeError("shared suh/svh differ across ranks") n = inter.shape[1] // 3 # coupled_scale_order gate_svh-up_svh-down_suh parts.append((inter[:, :n], inter[:, n:2 * n], inter[:, 2 * n:])) scales = CoupledScaleSet( gate_up_suh=suh.contiguous(), down_svh=svh.contiguous(), gate_svh=torch.cat([p[0] for p in parts], 1).contiguous(), up_svh=torch.cat([p[1] for p in parts], 1).contiguous(), down_suh=torch.cat([p[2] for p in parts], 1).contiguous(), source_path=R27, source_sha256="", source_metadata={}, tensor_hashes={}) scales.validate() return scales def routed_rows(cap, windows, expert): rows, weights = [], [] for w in windows: ids = np.asarray(cap.topk_ids[w * 2048:(w + 1) * 2048]) tok, slot = np.nonzero(ids == expert) if tok.size: rows.append(w * 2048 + tok.astype(np.int64)) weights.append(np.asarray(cap.topk_weights[w * 2048:(w + 1) * 2048])[tok, slot].astype(np.float32)) rows = np.concatenate(rows) hidden = torch.from_numpy(np.array(cap.hidden_words[rows], copy=True)).view(torch.bfloat16) return hidden, torch.from_numpy(np.concatenate(weights)) def hessian_chunks(x_cpu, route_cpu, fn, device, chunk=4096): """route_weighted_hessian (route^2, normalised) accumulated in chunks.""" total = None wsum = float(route_cpu.float().square().sum()) for s in range(0, x_cpu.shape[0], chunk): x = fn(x_cpu[s:s + chunk].to(device)) r = route_cpu[s:s + chunk].to(device).float() xw = x.float() * (r.square() / wsum).sqrt()[:, None] part = xw.t() @ xw total = part if total is None else total + part return total def mix(h_routed, h_all, beta): if beta == 0: return h_routed return ((1 - beta) * h_routed / torch.diagonal(h_routed).mean() + beta * h_all / torch.diagonal(h_all).mean()) def refit_cont(weight, rec, scales): blocks = weight.float().reshape(weight.shape[0], -1, 32) q = rec.reshape_as(blocks) / scales[..., None] num = (blocks.double() * q.double()).sum(-1) den = q.double().square().sum(-1) return torch.where(den > 0, (num / den.clamp_min(1e-30)).float(), scales).clamp_min(2.0 ** -126) def encode_v2(weight, hessian, bits, percdamp): lut = tm.state_lut(bits, alphabet="e4m3", law="mcg", compander_scale=2.0, device=weight.device) c_rms = float(lut.view(torch.float8_e4m3fn).float().square().mean().sqrt()) kw = dict(bits=bits, block_size=32, tailbite_context=128, percdamp=percdamp, column_block=128) scales = (weight.float().reshape(weight.shape[0], -1, 32).square().mean(-1).sqrt() / c_rms).clamp_min(2.0 ** -126) for _ in range(2): rec, _ = tm._encode_mxf_with_gptq_feedback(weight, hessian, scales, lut, **kw) scales = refit_cont(weight, rec, scales) scales = torch.pow(2.0, torch.round(torch.log2(scales))) for it in range(3): rec, states = tm._encode_mxf_with_gptq_feedback(weight, hessian, scales, lut, **kw) if it < 2: scales = tm._refit_power2_scales(weight, rec.reshape(weight.shape[0], -1, 32) / scales[..., None], scales) trellis = tn.pack_trellis_edges(states, bits) codes = tm.pack_ue8m0(scales) decoded = tm.decode_trellis_mxf(trellis, lut, codes, bits=bits, block_size=32, rows=weight.shape[0], width=weight.shape[1], device=weight.device) if not torch.equal(decoded, rec): raise RuntimeError("codec closure failed: packed payload does not decode to the encoder reconstruction") return rec, trellis.cpu().contiguous(), codes.cpu().contiguous() def main(): ap = argparse.ArgumentParser() ap.add_argument("--layer", type=int, required=True) ap.add_argument("--expert-start", type=int, default=0) ap.add_argument("--expert-end", type=int, default=288) ap.add_argument("--bits", type=int, required=True, choices=(3, 4, 5)) ap.add_argument("--beta", type=float, default=0.5) ap.add_argument("--percdamp", type=float, default=0.3) ap.add_argument("--out", type=Path, required=True) args = ap.parse_args() if args.out.exists(): raise FileExistsError(args.out) device = torch.device("cuda:0") torch.cuda.set_device(device) scales_all = r27_scales(args.layer) cap = LayerCapture(CAPTURE, args.layer, ROLES, max_samples=256, sample_offset=0, data_role="fit", sampling_strategy="domain-balanced") windows = list(cap.window_indices) assert len(windows) == 64 every = np.asarray([w * 2048 + t for w in windows for t in range(0, 2048, 4)], dtype=np.int64) hid_every = torch.from_numpy(np.array(cap.hidden_words[every], copy=True)).view(torch.bfloat16) ones_every = torch.ones(hid_every.shape[0]) source = IndexedCheckpoint(SOURCE, SOURCE / "model.safetensors.index.json") tensors, metrics, h_all_in = {}, [], None started = time.time() for expert in range(args.expert_start, args.expert_end): t0 = time.time() base = source.expert_prefix(args.layer, expert) orig = {p: source.get(f"{base}{p}.weight").to(device).float() for p in PROJ} sc = CoupledScaleSet( gate_up_suh=scales_all.gate_up_suh.to(device), gate_svh=scales_all.gate_svh[expert:expert + 1].to(device), up_svh=scales_all.up_svh[expert:expert + 1].to(device), down_suh=scales_all.down_suh[expert:expert + 1].to(device), down_svh=scales_all.down_svh.to(device), source_path=scales_all.source_path, source_sha256="", source_metadata={}, tensor_hashes={}) tw = encode_coupled_scale_weights(orig["gate_proj"], orig["up_proj"], orig["down_proj"], sc, expert=0, intermediate_draw=COUPLED_SIGN_DRAW) carrier = lambda h: coupled_input_carrier(h, sc.gate_up_suh, quantize=True) # noqa: E731 hid_r, rt_r = routed_rows(cap, windows, expert) if h_all_in is None and args.beta > 0: h_all_in = hessian_chunks(hid_every, ones_every, carrier, device) h_in = mix(hessian_chunks(hid_r, rt_r, carrier, device), h_all_in, args.beta) g_rec, g_tr, g_sc = encode_v2(tw[0], h_in, args.bits, args.percdamp) u_rec, u_tr, u_sc = encode_v2(tw[1], h_in, args.bits, args.percdamp) mid = lambda h: coupled_middle_carrier(carrier(h), g_rec, u_rec, sc, expert=0, # noqa: E731 intermediate_draw=COUPLED_SIGN_DRAW, quantize=True) h_dn = hessian_chunks(hid_r, rt_r, mid, device) if args.beta > 0: h_dn = mix(h_dn, hessian_chunks(hid_every, ones_every, mid, device), args.beta) d_rec, d_tr, d_sc = encode_v2(tw[2], h_dn, args.bits, args.percdamp) row = {"expert": expert, "routed_tokens": int(hid_r.shape[0]), "seconds": round(time.time() - t0, 2)} for name, rec, tr, codes, target in (("gate", g_rec, g_tr, g_sc, tw[0]), ("up", u_rec, u_tr, u_sc, tw[1]), ("down", d_rec, d_tr, d_sc, tw[2])): tensors[f"{expert}.{name}.trellis"] = tr tensors[f"{expert}.{name}.scale_ue8m0"] = codes row[f"{name}_tw_nmse"] = float((rec - target).double().square().sum() / target.double().square().sum()) metrics.append(row) print(json.dumps(row), flush=True) del orig, tw, h_in, h_dn, g_rec, u_rec, d_rec, hid_r, rt_r torch.cuda.empty_cache() meta = {"schema": "trellismx-codec-v2-chunk.v1", "layer": str(args.layer), "bits": str(args.bits), "expert_range": f"{args.expert_start}:{args.expert_end}", "beta": str(args.beta), "percdamp": str(args.percdamp), "encoder": Path(__file__).name, "encoder_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(), "calibration": "fit role, 64 windows: all routed tokens + every 4th token (all-token Hessian)"} args.out.parent.mkdir(parents=True, exist_ok=True) tmp = args.out.with_suffix(".partial") save_file(tensors, str(tmp), metadata=meta) tmp.rename(args.out) args.out.with_suffix(".metrics.json").write_text(json.dumps({"meta": meta, "rows": metrics, "seconds": time.time() - started}, indent=1)) print(json.dumps({"done": str(args.out), "seconds": round(time.time() - started, 1)}), flush=True) if __name__ == "__main__": main()