""" compress.py -- shared SVD-LLM compression pipeline (ONE code path for BASE & OURS). The arm identity lives entirely in the calib file (built by calib/build_calib.py). This module is arm-agnostic: same whitening math, same truncation, same pipeline. Faithful port of official SVD-LLM whitening (see docs/svdllm_port_diff.md): - single forward + forward hooks accumulate X^T X incrementally (O(d^2)/layer), no inps/outs buffering; - decomposition: torch.linalg.cholesky in fp64 with the official except branch (+(-lambda_min+1e-6) I on failure); eigh is a verified-equivalent backup, DEFAULT OFF (--decomp cholesky); - rank k = int(ratio * out * in / (out + in)); full SVD after whitening; - compresses the 224 block-internal target Linears; lm_head stays dense (guarded). Ratio convention (--model_ratio vs --ratio, exactly one required): --model_ratio 0.8 FULL-MODEL retention (方案B, LoRAP). The whole checkpoint is the denominator; embedding/lm_head are kept dense but COUNTED, so the per-layer ratio is derived DOWNWARD from the budget (common.rank_plan). This is what the main tables report and what SVD-LLM/LoRAP ratios mean. --ratio 0.7 LEGACY per-layer retention, embedding/lm_head outside the denominator. Retained only to reproduce archived results. Outputs (results/compress/ and the weights dir): per-layer A/B .pt + summary.json. """ import os import sys import time import argparse import torch import torch.nn as nn sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) import common as C # ── whitened truncation (official math). Reused by the equivalence self-check. ── def cholesky_factor(XtX, solve_dtype=torch.float64): """ Official Cholesky with except branch. Returns L (solve_dtype) s.t. L L^T = XtX. """ raw = XtX.to(solve_dtype) try: L = torch.linalg.cholesky(raw) except Exception: # official: shift by (-lambda_min + 1e-6) I then retry ev = torch.linalg.eigvalsh(raw) raw = raw + (-ev[0] + 1e-6) * torch.eye(raw.shape[0], dtype=solve_dtype, device=raw.device) L = torch.linalg.cholesky(raw) return L def eigh_factor(XtX, solve_dtype=torch.float64): """ Backup symmetric square root via eigh: L = V diag(sqrt(lambda)) s.t. L L^T = XtX. Different matrix than the Cholesky factor, but yields the SAME optimal rank-k W' because the truncation objective trace((W-W') XtX (W-W')^T) depends only on XtX. """ raw = XtX.to(solve_dtype) ev, V = torch.linalg.eigh(raw) ev = ev.clamp(min=0) L = V * ev.sqrt().unsqueeze(0) return L def whiten_truncate(W, XtX, ratio, decomp="cholesky", solve_dtype=torch.float64, out_dtype=torch.float32, k=None): """ Returns (A, B, k) with A@B the rank-k activation-aware approximation of W. Precision ladder matches official: factor/inv in solve_dtype (fp64), whitened SVD in out_dtype (fp32). `k` overrides the rank formula: under the full-model ratio convention the rank comes from the pre-computed plan (common.rank_plan), not from a scalar ratio. The truncation math is untouched either way -- only where k comes from changes. """ out_dim, in_dim = W.shape if decomp == "cholesky": L = cholesky_factor(XtX, solve_dtype) elif decomp == "eigh": L = eigh_factor(XtX, solve_dtype) else: raise ValueError(decomp) try: L_inv = torch.linalg.inv(L) except Exception: L = L + 1e-6 * torch.eye(L.shape[0], dtype=L.dtype, device=L.device) L_inv = torch.linalg.inv(L) Lf = L.to(out_dtype) Linv_f = L_inv.to(out_dtype) Wf = W.to(out_dtype) W_scale = Wf @ Lf U, S, Vt = torch.linalg.svd(W_scale, full_matrices=False) if k is None: k = C.rank_from_ratio(ratio, out_dim, in_dim) k = min(k, S.shape[0]) truc_s = S[:k] truc_u = U[:, :k] truc_v = Vt[:k, :] @ Linv_f sqrt_s = truc_s.sqrt() A = (truc_u * sqrt_s).contiguous() # out x k B = (torch.diag(sqrt_s) @ truc_v).contiguous() # k x in return A, B, k def asvd_truncate(W, absmean, ratio, alpha=0.5, out_dtype=torch.float32, k=None): """ ASVD (Yuan et al., hahnyuan/ASVD4LLM): DIAGONAL activation-aware scaling. scaling = (abs_mean_per_input_channel) ** alpha W_scale = W * scaling (per-column) -> SVD -> truncate rank k -> V /= scaling Same rank formula as SVD-LLM (equal param budget). Returns (A, B, k) with A@B the rank-k ASVD approximation of W. """ out_dim, in_dim = W.shape scaling = absmean.to(out_dtype).clamp(min=1e-6) ** alpha # (in_dim,) Wf = W.to(out_dtype) W_scale = Wf * scaling.view(1, -1) # W @ diag(scaling) U, S, Vt = torch.linalg.svd(W_scale, full_matrices=False) if k is None: k = C.rank_from_ratio(ratio, out_dim, in_dim) k = min(k, S.shape[0]) sqrt_s = S[:k].sqrt() A = (U[:, :k] * sqrt_s).contiguous() # out x k B = (torch.diag(sqrt_s) @ (Vt[:k, :] / scaling.view(1, -1))).contiguous() # k x in return A, B, k def collect_absmean_gpu(model, input_ids, target_mods, device, log_prefix=""): """ASVD statistic: per-input-channel mean |activation|, averaged over windows. Cheap (one vector/layer) so no binning; single forward pass, GPU-resident.""" acc, cnt = {}, {} def make_hook(name, in_dim): def hook(module, inp, out): x = inp[0].detach() am = x.abs().float().mean(dim=-2).reshape(-1) # (in_dim,) per window (batch=1) if name not in acc: acc[name] = torch.zeros(in_dim, dtype=torch.float32, device=device) cnt[name] = 0 acc[name] += am cnt[name] += 1 return hook handles = [m.register_forward_hook(make_hook(n, m.in_features)) for n, m in target_mods.items()] with torch.no_grad(): for i in range(input_ids.shape[0]): model(input_ids[i : i + 1].to(device)) for h in handles: h.remove() out = {n: (acc[n] / cnt[n]).cpu() for n in acc} torch.cuda.empty_cache() return out def collect_xtx_gpu(model, input_ids, target_mods, device, log_prefix=""): """ Accumulate X^T X for {name: module} in ONE forward pass, ON GPU (no per-forward D2H copy), offloading to CPU only at the end. Caller must ensure this group's total XtX (sum of in_features^2 * 4 bytes) fits on the GPU beside the model. Returns {name: cov_cpu_fp32 (already divided by token count)}. """ cov, ntok = {}, {} def make_hook(name, in_dim): def hook(module, inp, out): x = inp[0].detach().reshape(-1, in_dim).float() if name not in cov: cov[name] = torch.zeros(in_dim, in_dim, dtype=torch.float32, device=device) ntok[name] = 0 cov[name].addmm_(x.t(), x) # cov += x^T x, in-place on GPU ntok[name] += x.shape[0] return hook handles = [m.register_forward_hook(make_hook(n, m.in_features)) for n, m in target_mods.items()] t0 = time.time() with torch.no_grad(): for i in range(input_ids.shape[0]): model(input_ids[i : i + 1].to(device)) if i % 200 == 0: print(f" {log_prefix}collect {i}/{input_ids.shape[0]} " f"({time.time()-t0:.0f}s)", flush=True) for h in handles: h.remove() out = {} for n in list(cov): out[n] = (cov[n] / ntok[n]).cpu() cov[n] = None del cov torch.cuda.empty_cache() return out def plan_bins(targets, budget_bytes): """Greedy-bin target linears so each bin's XtX total <= budget (fits on GPU).""" bins, cur, cur_b = [], {}, 0 for name, mod in targets: b = mod.in_features * mod.in_features * 4 if cur and cur_b + b > budget_bytes: bins.append(cur) cur, cur_b = {}, 0 cur[name] = mod cur_b += b if cur: bins.append(cur) return bins def main(): ap = argparse.ArgumentParser() ap.add_argument("--calib", required=True, help="calib .pt from build_calib.py") # Exactly one ratio convention must be named -- there is no default, because # the two mean different compression strengths and silently picking one is how # the old layer-wise numbers became incomparable to the literature. g = ap.add_mutually_exclusive_group(required=True) g.add_argument("--model_ratio", type=float, default=None, help="FULL-MODEL parameter retention (方案B, LoRAP convention). " "Denominator = the whole checkpoint incl. embedding+lm_head. " "Per-layer ranks are derived from this budget; the resulting " "layer ratio is LOWER (= stronger compression) than this " "number. USE THIS for the main tables.") g.add_argument("--ratio", type=float, default=None, help="LEGACY per-layer retention: applied directly to every target " "Linear, embedding/lm_head excluded from the denominator. " "Kept only to reproduce the archived layer-wise results; " "NOT comparable to SVD-LLM/LoRAP reported ratios.") ap.add_argument("--layer_type", choices=["all", "attn", "mlp"], default="all") ap.add_argument("--decomp", choices=["cholesky", "eigh", "asvd"], default="cholesky") ap.add_argument("--alpha", type=float, default=0.5, help="ASVD scaling exponent (only used when --decomp asvd)") ap.add_argument("--model_path", type=str, default=C.DEFAULT_MODEL_PATH) ap.add_argument("--save_path", required=True) ap.add_argument("--batch_size", type=int, default=1) ap.add_argument("--nsamples", type=int, default=0, help="cap calibration to the first N windows (0 = use all)") ap.add_argument("--xtx_budget_gb", type=float, default=0.0, help="max GB of XtX kept on GPU per collection pass. " "0 = AUTO: measure free GPU memory after the model is " "resident and keep an activation headroom margin, so the " "same command runs on a 24GB or a 48GB card (only the " "number of collection passes changes -- the result is " "bit-identical either way, binning never touches math). " "The old fixed 20GB default assumed a 46GB card and OOMs " "on a 24GB A5000, where only ~8GB is free beside the " "16GB bf16 model.") ap.add_argument("--xtx_headroom_gb", type=float, default=6.0, help="GPU memory reserved for activations/workspace in AUTO mode") args = ap.parse_args() device = "cuda" if torch.cuda.is_available() else "cpu" ghash = C.git_hash() os.makedirs(args.save_path, exist_ok=True) t0 = time.time() # ── load calib (carries the arm identity) ───────────────────────────────── blob = torch.load(args.calib, map_location="cpu") input_ids = blob["input_ids"] if args.nsamples and args.nsamples < input_ids.shape[0]: input_ids = input_ids[: args.nsamples] arm = blob.get("arm", "unknown") ratio_mode = "model" if args.model_ratio is not None else "layer" print(f"[compress] arm={arm} calib={os.path.basename(args.calib)} " f"N={input_ids.shape[0]} ratio_mode={ratio_mode} " f"model_ratio={args.model_ratio} ratio={args.ratio} " f"layer_type={args.layer_type} decomp={args.decomp}") # ── load model ──────────────────────────────────────────────────────────── model, tokenizer = C.load_model(args.model_path, device=device) n_targets = C.count_target_linears(model, args.layer_type) print(f"[compress] target linears ({args.layer_type}) = {n_targets}") # ── rank allocation ─────────────────────────────────────────────────────── # Derived from the ARCHITECTURE ONLY (dims + full-model param census), never # from the calibration data, so BASE and OURS are rank-identical by # construction rather than by luck. The effective-ratio gate fires here, # BEFORE the hours of XtX collection, so an allocation bug costs seconds. if ratio_mode == "model": plan = C.rank_plan(model, args.model_ratio, args.layer_type) C.assert_effective_ratio(plan) ranks = plan["ranks"] print(f"[plan] FULL-MODEL ratio convention (方案B / LoRAP)\n" f" Param_total = {plan['param_total']:,}\n" f" P_fixed = {plan['p_fixed']:,} " f"({plan['p_fixed_share']:.2%}, embedding+lm_head+norms+biases, uncompressed)\n" f" P_targets = {plan['p_targets']:,} ({plan['n_targets']} Linears)\n" f" target budget = {plan['target_budget_params']:,}\n" f" layer_ratio = {plan['layer_ratio']:.6f} " f"(= {1-plan['layer_ratio']:.2%} removed per target layer)\n" f" post-compress = {plan['post_compression_total_params']:,}\n" f" effective full-model ratio = " f"{plan['effective_model_ratio']:.6f} (nominal {args.model_ratio})\n" f" rank_plan_hash = {plan['rank_plan_hash']}") else: plan = None ranks = None print(f"[plan] LEGACY per-layer ratio convention: ratio={args.ratio} applied " f"to each target Linear; embedding/lm_head NOT in the denominator. " f"These numbers are not comparable to SVD-LLM/LoRAP model-level ratios.") # ── X^T X collection: GPU-resident accumulation, memory-binned ──────────── # XtX stays on GPU (no per-forward D2H copy) and is offloaded once per bin. # all-linear XtX (~32GB) exceeds a 46GB GPU beside the 16GB model, so targets # are split into bins that each fit; one forward pass per bin. targets = list(C.iter_target_linears(model, args.layer_type)) t_collect = time.time() if args.decomp == "asvd": # ASVD: per-channel abs-mean (tiny vectors) -> one pass, no binning. cov = collect_absmean_gpu(model, input_ids, dict(targets), device) print(f"[compress] ASVD abs-mean collected for {len(cov)} layers " f"({time.time()-t_collect:.0f}s)") else: # XtX stays on GPU (no per-forward D2H copy), offloaded once per bin. # LLaDA's all-linear XtX is ~32GB (the block-internal ff_out alone: # 12288^2*4B = 604MB x 32), far beyond any single GPU beside the 16GB bf16 # model, so targets are split into bins that each fit; one forward pass per # bin. Binning is purely a memory schedule -- every layer still sees the # full calibration set, so the result is bit-identical at any budget. # AUTO sizing (default) reads the free memory AFTER the model is resident, # so the same command survives both a 24GB A5000 and a 48GB L40S. budget_gb = args.xtx_budget_gb if budget_gb <= 0: if device == "cuda": free_b, total_b = torch.cuda.mem_get_info() budget_gb = max(2.0, free_b / 1024 ** 3 - args.xtx_headroom_gb) print(f"[compress] AUTO xtx budget: free={free_b/1024**3:.1f}GB " f"total={total_b/1024**3:.1f}GB headroom={args.xtx_headroom_gb}GB " f"-> budget={budget_gb:.1f}GB") else: budget_gb = 8.0 bins = plan_bins(targets, budget_gb * 1024 ** 3) print(f"[compress] {len(targets)} target linears -> {len(bins)} XtX pass(es) " f"(budget {budget_gb:.1f}GB/pass)") cov = {} for bi, b in enumerate(bins): gb = sum(m.in_features ** 2 * 4 for m in b.values()) / 1024 ** 3 print(f"[compress] XtX pass {bi+1}/{len(bins)}: {len(b)} layers, {gb:.1f}GB", flush=True) cov.update(collect_xtx_gpu(model, input_ids, b, device, log_prefix=f"p{bi+1} ")) print(f"[compress] XtX collected for {len(cov)} layers " f"({time.time()-t_collect:.0f}s)") # ── whiten + truncate + replace ─────────────────────────────────────────── t_dec = time.time() summary_layers = {} orig_params = comp_params = n_comp = n_skip = skip_params = 0 for name, mod in list(C.iter_target_linears(model, args.layer_type)): if name not in cov: continue W = mod.weight.data.cpu().float() out_dim, in_dim = W.shape if ranks is not None: k = ranks.get(name) # from the pre-computed full-model plan if k is None: # plan already decided: no saving -> dense n_skip += 1 skip_params += out_dim * in_dim continue else: k = C.rank_from_ratio(args.ratio, out_dim, in_dim) if not C.is_compression_beneficial(k, out_dim, in_dim): n_skip += 1 skip_params += out_dim * in_dim continue if args.decomp == "asvd": A, B, k = asvd_truncate(W, cov[name], args.ratio, alpha=args.alpha, k=k) else: A, B, k = whiten_truncate(W, cov[name], args.ratio, decomp=args.decomp, k=k) prefix = name.replace(".", "_") torch.save(A.cpu(), os.path.join(args.save_path, f"{prefix}_A.pt")) torch.save(B.cpu(), os.path.join(args.save_path, f"{prefix}_B.pt")) parent, attr = C.get_parent_attr(model, name) # A/B are already saved to disk; the compressed layer is NOT needed on GPU # (no PPL eval follows). Keep the replacement on CPU so the GPU stays flat at # ~model size during the SVD phase instead of accumulating 224 layers -> OOM. setattr(parent, attr, C.LowRankLinear(A, B, bias=None)) orig_params += out_dim * in_dim comp_params += out_dim * k + k * in_dim n_comp += 1 summary_layers[name] = {"out": out_dim, "in": in_dim, "k": int(k)} # memory hygiene: drop refs, do NOT move 8B params to CPU before del del W, A, B torch.cuda.empty_cache() C.assert_head_dense(model) # lm_head hard guard AFTER replacement print(f"[compress] compressed={n_comp} skipped={n_skip} " f"({time.time()-t_dec:.0f}s)") # ── measured full-model ratio (not predicted: counted off the real weights) ── # Closes the loop on the plan: what the plan promised must equal what the 224 # replacements actually cost. Any mismatch means the plan and the executed # compression disagree, which would make the reported ratio a fiction. measured = None if plan is not None: realized_targets = comp_params + skip_params if realized_targets != plan["realized_target_params"]: raise RuntimeError( f"PLAN/EXECUTION MISMATCH: target params realized={realized_targets:,} " f"but plan predicted {plan['realized_target_params']:,}" ) measured = (plan["p_fixed"] + realized_targets) / plan["param_total"] C.assert_effective_ratio({**plan, "effective_model_ratio": measured}) print(f"[compress] MEASURED full-model ratio = {measured:.6f} " f"(nominal {args.model_ratio}) -- verified against the plan") # ── summary ─────────────────────────────────────────────────────────────── root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) peak = C.peak_rss_gb() summary = { "arm": arm, "model_id": C.DEFAULT_MODEL_ID, "calib_file": os.path.basename(args.calib), "calib_git_hash": blob.get("git_hash"), "ratio_mode": ratio_mode, "model_ratio": args.model_ratio, "ratio": args.ratio, "layer_ratio": (plan["layer_ratio"] if plan else args.ratio), "layer_type": args.layer_type, "decomp": args.decomp, "model_path": args.model_path, "n_target_linears": n_targets, "n_compressed": n_comp, "n_skipped": n_skip, "orig_params": orig_params, "compressed_params": comp_params, "skipped_params": skip_params, "kept_fraction": (comp_params / orig_params) if orig_params else None, # ── full-model ratio bookkeeping (the number the paper reports) ────────── "param_total": (plan["param_total"] if plan else None), "p_fixed": (plan["p_fixed"] if plan else None), "p_targets": (plan["p_targets"] if plan else None), "p_fixed_share": (plan["p_fixed_share"] if plan else None), "target_budget_params": (plan["target_budget_params"] if plan else None), "post_compression_total_params": ( plan["p_fixed"] + comp_params + skip_params if plan else None), "effective_model_ratio": measured, "rank_plan_hash": (plan["rank_plan_hash"] if plan else None), "ratio_semantics": ( "ratio_mode=model (方案B/LoRAP): model_ratio = FULL-MODEL parameter " "RETENTION, denominator = every param incl. embedding+lm_head. " "layer_ratio is DERIVED (= (model_ratio*Param_total - P_fixed)/P_targets) " "and is LOWER than model_ratio, i.e. target layers are compressed HARDER " "to pay for the untouched embedding/lm_head. effective_model_ratio is the " "MEASURED full-model retention -- report that number, not any symbol. " "kept_fraction remains the target-only fraction (legacy column)." if plan else "ratio_mode=layer (LEGACY): ratio = per-target-Linear RETENTION fraction " "(k=int(ratio*out*in/(out+in))); embedding/lm_head are NOT in the " "denominator, so this is NOT comparable to SVD-LLM/LoRAP model-level " "ratios. kept_fraction is the MEASURED compressed/orig over targets only."), "per_layer": summary_layers, "git_hash": ghash, "peak_rss_gb": peak, "wall_clock_sec": time.time() - t0, "save_path": args.save_path, } # written both next to weights and into results/compress/ C.dump_json(summary, os.path.join(args.save_path, "compression_summary.json")) # ratio tag names the CONVENTION too, so a model-ratio artifact can never be # mistaken for a legacy layer-ratio one on the filesystem. rtag = (f"m{args.model_ratio}" if ratio_mode == "model" else f"r{args.ratio}") res_name = f"{ghash}_{arm}_{args.layer_type}_{rtag}_summary.json" C.dump_json(summary, os.path.join(root, "results", "compress", res_name)) print(f"[compress] summary -> {args.save_path}/compression_summary.json") print(f"[compress] peak_rss={peak:.1f}GB wall={ (time.time()-t0)/60:.1f}min " f"kept_targets={summary['kept_fraction']:.6f} " f"effective_model_ratio={measured if measured is None else f'{measured:.6f}'}") if __name__ == "__main__": main()