| """Sarvam-Prism -> INT4 package for <=16 GB GPUs: routed experts to INT4 (GPTQ, group 128, symmetric), everything else kept as shipped. |
| |
| Stages (one process; calibration activations live in RAM): |
| 1. load the original FP8 model on the big GPU, run the evaluation suite -> ref.json |
| 2. run calibration requests, record each MoE layer's input tokens (and their routing) |
| 3. free the GPU; for each MoE layer: build per-expert Hessians, GPTQ all 128 experts batched on the GPU, |
| pack to int4, write the output layer file, delete the source layer file (space on /dev/shm is limited) |
| """ |
| import gc, json, os, shutil, sys, time |
| import torch |
| import torch.nn.functional as F |
| from safetensors import safe_open |
| from safetensors.torch import save_file |
|
|
| SRC, OUT, WORK = "/dev/shm/src", "/dev/shm/out", "/workspace/work" |
| GROUP, QMAX, QMIN = 128, 7, -8 |
| CALIB_TOKENS = int(os.environ.get("CALIB_TOKENS", 262144)) |
| CLIPS = torch.tensor([1.0, 0.96, 0.92, 0.88, 0.84, 0.80]) |
| os.makedirs(OUT, exist_ok=True); os.makedirs(WORK, exist_ok=True) |
| log = lambda *a: print(time.strftime("%H:%M:%S"), *a, flush=True) |
|
|
|
|
| |
| def reference_and_calibrate(): |
| sys.path.insert(0, SRC); sys.path.insert(0, os.path.dirname(__file__)); sys.path.insert(0, WORK) |
| import prism |
| from suite import run |
| from calib_requests import make_requests |
| t = time.time() |
| model = prism.Prism.from_pretrained(SRC, cuda_graphs=False) |
| log(f"reference loaded in {time.time()-t:.0f}s, GPU mem {torch.cuda.memory_allocated()/1e9:.1f} GB") |
| if not os.path.exists(f"{WORK}/ref.json"): |
| run(model, f"{WORK}/ref.json"); log("reference suite done") |
| store, state = {}, {"on": False} |
| for name, m in model.backbone.named_modules(): |
| if type(m).__name__ == "SarvamMLAMoE": |
| li = int(name.split(".")[1]); orig = m._moe_dispatch; store[li] = ([], []) |
| def wrap(x, ids, w, li=li, orig=orig): |
| if state["on"]: |
| lens = prism.SPLIT["len"]; B = len(lens); L = x.shape[0] // B |
| sel = torch.cat([torch.arange(b * L, b * L + n, device=x.device) for b, n in enumerate(lens)]) |
| store[li][0].append(x[sel].to(torch.bfloat16).cpu()); store[li][1].append(ids[sel].to(torch.uint8).cpu()) |
| return orig(x, ids, w) |
| m._moe_dispatch = wrap |
| reqs = make_requests(os.path.join(os.path.dirname(__file__), "wiki.json")) |
| state["on"] = True; n_tok, t = 0, time.time() |
| for i, (s, q) in enumerate(reqs): |
| model.decide(state=s, questions=q) |
| n_tok = sum(x.shape[0] for x in store[1][0]) |
| if i % 200 == 0: log(f"calib req {i} tokens {n_tok}") |
| if n_tok >= CALIB_TOKENS: break |
| state["on"] = False |
| log(f"calibration: {i+1} requests, {n_tok} tokens, {time.time()-t:.0f}s") |
| calib = {li: (torch.cat(xs), torch.cat(ids)) for li, (xs, ids) in store.items()} |
| del model, store; gc.collect(); torch.cuda.empty_cache() |
| return calib |
|
|
|
|
| |
| def group_scale(w): |
| """w: [E, N, G] -> per-(E,N) scale chosen from CLIPS by minimum rounding MSE.""" |
| amax = w.abs().amax(-1).clamp_min(1e-8) |
| best, best_err = None, None |
| for c in CLIPS.tolist(): |
| s = amax * c / QMAX |
| err = ((torch.clamp(torch.round(w / s[..., None]), QMIN, QMAX) * s[..., None] - w) ** 2).sum(-1) |
| if best is None: |
| best, best_err = s, err |
| else: |
| m = err < best_err; best = torch.where(m, s, best); best_err = torch.where(m, err, best_err) |
| return best |
|
|
|
|
| @torch.no_grad() |
| def gptq(W, H, cnt, damp=0.01, min_tokens=512): |
| """Batched GPTQ. W: [E, N, K] fp32, H: [E, K, K] fp32 (sum of x x^T), cnt: calibration tokens per expert. |
| Experts seen by fewer than min_tokens tokens get H = I (plain round-to-nearest with clip search: no error |
| feedback fitted to a handful of samples). Inputs never seen (zero diagonal) keep their weights (diag set to the |
| mean) so rare inputs at deployment are still served. Returns int q [E,N,K] and fp32 scales [E,N,K/G].""" |
| E, N, K = W.shape |
| W = W.clone(); H = H.clone() |
| few = (cnt < min_tokens).to(H.device) |
| if few.any(): |
| H[few] = torch.eye(K, device=H.device) |
| diag = torch.diagonal(H, dim1=1, dim2=2) |
| mean = diag.mean(-1, keepdim=True).clamp_min(1e-8) |
| dead = diag <= 0 |
| diag.copy_(torch.where(dead, mean.expand_as(diag), diag)) |
| diag += damp * mean |
| L, info = torch.linalg.cholesky_ex(H) |
| if info.any(): |
| bad = info.nonzero().flatten() |
| for e in bad.tolist(): |
| d = torch.diagonal(H[e]); d += 0.1 * d.mean() |
| L, info = torch.linalg.cholesky_ex(H) |
| assert not info.any(), "Cholesky failed after extra damping" |
| Hinv = torch.cholesky_inverse(L) |
| Hinv = torch.linalg.cholesky(Hinv, upper=True) |
| Q = torch.empty(E, N, K, dtype=torch.int8, device=W.device) |
| S = torch.empty(E, N, K // GROUP, dtype=torch.float32, device=W.device) |
| for g, i1 in enumerate(range(0, K, GROUP)): |
| i2 = i1 + GROUP |
| W1 = W[:, :, i1:i2].clone(); Err = torch.empty_like(W1); Hi = Hinv[:, i1:i2, i1:i2] |
| s = group_scale(W1); S[:, :, g] = s |
| for i in range(GROUP): |
| w = W1[:, :, i] |
| q = torch.clamp(torch.round(w / s), QMIN, QMAX) |
| Q[:, :, i1 + i] = q.to(torch.int8) |
| err = (w - q * s) / Hi[:, i, i][:, None] |
| W1[:, :, i:] -= err[:, :, None] * Hi[:, i, i:][:, None, :] |
| Err[:, :, i] = err |
| if i2 < K: |
| W[:, :, i2:] -= torch.bmm(Err, Hinv[:, i1:i2, i2:]) |
| return Q, S |
|
|
|
|
| def rtn(W): |
| E, N, K = W.shape |
| Wg = W.view(E, N, K // GROUP, GROUP) |
| s = Wg.abs().amax(-1).clamp_min(1e-8) / QMAX |
| return torch.clamp(torch.round(Wg / s[..., None]), QMIN, QMAX).view(E, N, K).to(torch.int8), s |
|
|
|
|
| def dequant(Q, S): |
| E, N, K = Q.shape |
| return (Q.float().view(E, N, K // GROUP, GROUP) * S[..., None]).view(E, N, K) |
|
|
|
|
| def pack(Q): |
| """int q in [-8,7], [E,N,K] -> uint8 [E,N,K/2]. Within each 128-group, byte j holds element j (low nibble) and |
| element j+64 (high nibble).""" |
| E, N, K = Q.shape |
| u = (Q.to(torch.int16) + 8).to(torch.uint8).view(E, N, K // GROUP, 2, GROUP // 2) |
| return (u[:, :, :, 0] | (u[:, :, :, 1] << 4)).reshape(E, N, K // 2).contiguous() |
|
|
|
|
| def unpack(P): |
| E, N, K2 = P.shape |
| p = P.view(E, N, K2 // (GROUP // 2), GROUP // 2) |
| lo, hi = (p & 15).to(torch.int8) - 8, (p >> 4).to(torch.int8) - 8 |
| return torch.stack([lo, hi], 3).reshape(E, N, K2 * 2) |
|
|
|
|
| def hessians(X, ids, n_exp, fn=None, dim=None): |
| """Per-expert sum x x^T over the tokens routed to it. fn(e, x) maps the token inputs (e.g. to the down-proj |
| input). Returns H [E, D, D] fp32 and token counts.""" |
| flat = ids.long() |
| D = dim or X.shape[1] |
| H = torch.zeros(n_exp, D, D, device="cuda") |
| cnt = torch.zeros(n_exp, dtype=torch.long) |
| tok = torch.arange(X.shape[0], device=X.device).repeat_interleave(ids.shape[1]) |
| order = flat.flatten().argsort() |
| e_sorted, t_sorted = flat.flatten()[order], tok[order] |
| bounds = torch.searchsorted(e_sorted, torch.arange(n_exp + 1, device=X.device)).tolist() |
| for e in range(n_exp): |
| t = t_sorted[bounds[e]:bounds[e + 1]] |
| cnt[e] = t.numel() |
| for c in range(0, t.numel(), 32768): |
| x = X[t[c:c + 32768]].float() |
| if fn is not None: x = fn(e, x) |
| H[e].addmm_(x.T, x) |
| return H, cnt |
|
|
|
|
| def rel_err(Wq, W, H): |
| """Hessian-weighted relative output error: tr(D H D^T) / tr(W H W^T), averaged over experts.""" |
| D = Wq - W |
| num = torch.einsum("enk,ekj,enj->e", D, H, D); den = torch.einsum("enk,ekj,enj->e", W, H, W) |
| return (num / den.clamp_min(1e-12)).mean().item() |
|
|
|
|
| |
| def quantize_layer(li, X, ids, stats): |
| f = f"model-layer-{li:02d}.safetensors" |
| tensors = {} |
| with safe_open(f"{SRC}/{f}", "pt", device="cuda") as sf: |
| for k in sf.keys(): tensors[k] = sf.get_tensor(k) |
| p = f"model.layers.{li}.mlp.experts." |
| gu = tensors.pop(p + "gate_up_proj").float() * tensors.pop(p + "gate_up_proj_scale")[..., None] |
| dn = tensors.pop(p + "down_proj").float() * tensors.pop(p + "down_proj_scale")[..., None] |
| E, I2, Hd = gu.shape; I = I2 // 2 |
| X, ids = X.cuda(), ids.cuda() |
| t = time.time() |
| Hgu, cnt = hessians(X, ids, E) |
| Qgu, Sgu = gptq(gu, Hgu, cnt) |
| guq = dequant(Qgu, Sgu) |
| e_gu = rel_err(guq, gu, Hgu); e_gu_rtn = rel_err(dequant(*rtn(gu)), gu, Hgu); del Hgu |
| def act(e, x): |
| h = x @ guq[e].T |
| return F.silu(h[:, :I]) * h[:, I:] |
| Hdn, _ = hessians(X, ids, E, fn=act, dim=I) |
| Qdn, Sdn = gptq(dn, Hdn, cnt) |
| e_dn = rel_err(dequant(Qdn, Sdn), dn, Hdn); e_dn_rtn = rel_err(dequant(*rtn(dn)), dn, Hdn); del Hdn |
| assert torch.equal(unpack(pack(Qgu)), Qgu) and torch.equal(unpack(pack(Qdn)), Qdn) |
| tensors[p + "gate_up_qweight"] = pack(Qgu).cpu(); tensors[p + "gate_up_scales"] = Sgu.float().cpu() |
| tensors[p + "down_qweight"] = pack(Qdn).cpu(); tensors[p + "down_scales"] = Sdn.float().cpu() |
| assert torch.isfinite(Sgu).all() and (Sgu > 0).all() and torch.isfinite(Sdn).all() and (Sdn > 0).all() |
| tensors = {k: v.cpu().contiguous() for k, v in tensors.items()} |
| save_file(tensors, f"{OUT}/{f}.tmp", metadata={"format": "pt"}) |
| os.replace(f"{OUT}/{f}.tmp", f"{OUT}/{f}") |
| stats[li] = dict(gate_up_gptq=e_gu, gate_up_rtn=e_gu_rtn, down_gptq=e_dn, down_rtn=e_dn_rtn, |
| min_tokens=int(cnt.min()), rtn_experts=int((cnt < 512).sum()), median_tokens=int(cnt.median()), seconds=round(time.time() - t, 1)) |
| log(f"layer {li}: {stats[li]}") |
| os.remove(f"{SRC}/{f}") |
| return list(tensors) |
|
|
|
|
| def main(): |
| stats_path = f"{WORK}/quant_stats.json" |
| stats = json.load(open(stats_path)) if os.path.exists(stats_path) else {} |
| stats = {int(k): v for k, v in stats.items()} |
| calib = reference_and_calibrate() |
| failed = [] |
| cfg = json.load(open(f"{SRC}/config.json")) |
| for li in range(cfg["first_k_dense_replace"], cfg["num_hidden_layers"]): |
| if os.path.exists(f"{OUT}/model-layer-{li:02d}.safetensors"): |
| continue |
| X, ids = calib.pop(li) |
| try: |
| quantize_layer(li, X, ids, stats) |
| except Exception: |
| import traceback; traceback.print_exc() |
| log(f"layer {li} FAILED; continuing (rerun to retry)"); failed.append(li) |
| json.dump(stats, open(stats_path, "w"), indent=1) |
| del X, ids; gc.collect(); torch.cuda.empty_cache() |
| if failed: |
| log(f"FAILED layers: {failed}"); sys.exit(1) |
| |
| for f in os.listdir(SRC): |
| if f.startswith(".") or f in ("prism.py", "example.py", "requirements.txt", "README.md", "model.safetensors.index.json"): |
| continue |
| if not os.path.exists(f"{OUT}/{f}"): |
| (shutil.copytree if os.path.isdir(f"{SRC}/{f}") else shutil.copy2)(f"{SRC}/{f}", f"{OUT}/{f}") |
| wm = {} |
| for f in sorted(os.listdir(OUT)): |
| if f.endswith(".safetensors") and f.startswith("model-"): |
| with safe_open(f"{OUT}/{f}", "pt") as sf: |
| for k in sf.keys(): wm[k] = f |
| json.dump({"metadata": {}, "weight_map": wm}, open(f"{OUT}/model.safetensors.index.json", "w"), indent=1) |
| cfg["quantization"] = { |
| "routed_experts": {"format": "int4", "method": "GPTQ (damp 0.01, MSE clip search)", "symmetric": True, |
| "group_size": GROUP, "scale_dtype": "float32", |
| "packing": "uint8 [E, N, K/2]; in each 128-group byte j = elem j (low nibble) | elem j+64 (high nibble); value = nibble - 8", |
| "calibration": f"{CALIB_TOKENS} tokens of Prism-format requests over multilingual Wikipedia (11 languages)"}, |
| "other_projections": {"format": "float8_e4m3fn", "weight_scale": "per output row (absmax / 448)", "note": "unchanged from sarvamai/prism"}, |
| "embeddings_norms": "bfloat16", "router": "float32"} |
| json.dump(cfg, open(f"{OUT}/config.json", "w"), indent=2) |
| log("done") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
