Blame for code/compress/compress.py

"""
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()