weight_rank.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
# This file is part of the Aquin Engine. Unauthorized copying, modification,
# distribution, or use of this file, via any medium, is strictly prohibited.
# Proprietary and confidential. See LICENSE for terms.
 
"""Stable rank analysis for transformer weight matrices."""
from __future__ import annotations
 
import re
 
import torch
 
 
_LAYER_RE = re.compile(r"(?:^|\.)(?:layers?|layer|block)\.?(\d+)(?:\.|$)", re.I)
_EXPERT_WEIGHT_RE = re.compile(r"\.experts?\.\d+", re.I)
 
 
def _is_heavy_weight_model(model) -> bool:
    """Large / MoE models fill VRAM — weight analysis must run on CPU and skip experts."""
    cfg = getattr(model, "cfg", None)
    d_model = int(getattr(cfg, "d_model", 0) or 0)
    if d_model >= 4096:
        return True
    n_layers = int(getattr(cfg, "n_layers", 0) or 0)
    if d_model >= 2048 and n_layers >= 24:
        return True
    # Avoid iterating every MoE expert parameter — that can hang or exhaust host RAM.
    return False
 
 
def _should_skip_weight_param(name: str, *, heavy: bool) -> bool:
    """Skip tensors that are redundant or too large for in-GPU analysis."""
    if not heavy:
        return False
    kl = name.lower()
    if _EXPERT_WEIGHT_RE.search(name):
        return True
    if "word_embeddings" in kl or kl.endswith("lm_head.weight"):
        return True
    if "router" in kl or "gate" in kl and "mlp" in kl:
        return False  # keep MoE router/gate — small and informative
    return False
 
 
def _layer_from_key(key: str) -> int:
    m = _LAYER_RE.search(key.replace("_", "."))
    return int(m.group(1)) if m else -1
 
 
def _hf_matrix_label(key: str) -> str:
    kl = key.lower()
    if "self_attn" in kl or ".attn." in kl:
        for tag, label in (
            ("q_proj", "Q"),
            ("k_proj", "K"),
            ("v_proj", "V"),
            ("o_proj", "O"),
            ("out_proj", "O"),
        ):
            if tag in kl:
                return label
    if "mlp" in kl or "feed_forward" in kl:
        if "gate" in kl:
            return "MLP_gate"
        if "down" in kl:
            return "MLP_down"
        if "up" in kl:
            return "MLP_up"
        return "MLP"
    if "conv" in kl or "shortconv" in kl:
        if "in_proj" in kl:
            return "Conv_in"
        if "out_proj" in kl:
            return "Conv_out"
        if re.search(r"\.conv\.|conv\d*\.weight$", kl):
            return "Conv_dw"
    for tag, label in (
        ("q_proj", "Q"),
        ("k_proj", "K"),
        ("v_proj", "V"),
        ("o_proj", "O"),
        ("out_proj", "O"),
        ("gate_proj", "MLP_gate"),
        ("up_proj", "MLP_up"),
        ("down_proj", "MLP_down"),
        ("query_key_value", "QKV"),
        ("c_attn", "Attn"),
        ("c_fc", "MLP_up"),
        ("c_proj", "MLP_down"),
    ):
        if tag in kl:
            return label
    parts = key.split(".")
    return parts[-2] if len(parts) >= 2 else parts[-1]
 
 
def _weight_root(model) -> object:
    if hasattr(model, "hf_model"):
        return model.hf_model
    return model
 
 
def _is_hf_weight_model(model) -> bool:
    if hasattr(model, "hf_model"):
        return True
    if not hasattr(model, "blocks"):
        return True
    try:
        block = model.blocks[0]
        return not (hasattr(block, "attn") and hasattr(block.attn, "W_Q"))
    except (IndexError, TypeError, AttributeError):
        return True
 
 
def _scalar(t: torch.Tensor) -> float:
    return t.detach().float().reshape(-1)[0].item()
 
 
def _to_2d(w: torch.Tensor) -> torch.Tensor:
    """HookedTransformer attn weights are often 3D (n_heads, d_in, d_out)."""
    f = w.detach().float()
    if f.ndim == 1:
        return f.unsqueeze(0)
    if f.ndim == 2:
        return f
    return f.reshape(-1, f.shape[-1])
 
 
def _matrix_ranks(w: torch.Tensor) -> tuple[float, float]:
    # SVD on CPU — model may already occupy nearly all VRAM (e.g. Sarvam 30B MoE).
    f = _to_2d(w).detach().float().cpu()
    if f.numel() < 2:
        return 1.0, 1.0
    try:
        sv = torch.linalg.svdvals(f)
    except Exception:
        return 0.0, 0.0
    if sv.numel() == 0:
        return 0.0, 0.0
    s0 = _scalar(sv.reshape(-1)[0])
    if s0 < 1e-12:
        return 0.0, 0.0
    ssum = _scalar(sv.sum())
    s2sum = _scalar((sv ** 2).sum())
    stable = (ssum ** 2) / max(s2sum, 1e-12)
    nuclear_ratio = ssum / (s0 * sv.numel())
    return round(stable, 4), round(nuclear_ratio, 4)
 
 
def run_weight_rank_hf(model, collapse_threshold: float = 0.1) -> dict:
    """Stable rank over HuggingFace parameter names (LFM, Sarvam, etc.)."""
    root = _weight_root(model)
    cfg = getattr(model, "cfg", None)
    n_layers = int(getattr(cfg, "n_layers", 0) or 0)
    heavy = _is_heavy_weight_model(model)
    matrices: list[dict] = []
    seen_layers: set[int] = set()
 
    for name, param in root.named_parameters():
        if param.dim() < 2 or param.numel() < 256:
            continue
        if _should_skip_weight_param(name, heavy=heavy):
            continue
        layer = _layer_from_key(name)
        if layer < 0:
            continue
        seen_layers.add(layer)
        matrix = _hf_matrix_label(name)
        stable, nuclear = _matrix_ranks(param)
        matrices.append({
            "layer": layer,
            "matrix": matrix,
            "module": name,
            "stable_rank": stable,
            "nuclear_norm_ratio": nuclear,
            "collapsed": stable < collapse_threshold,
        })
 
    if not n_layers and seen_layers:
        n_layers = max(seen_layers) + 1
 
    collapsed_count = sum(1 for m in matrices if m["collapsed"])
    return {
        "n_layers": n_layers,
        "matrices": matrices,
        "collapse_threshold": collapse_threshold,
        "collapsed_count": collapsed_count,
        "backend": "hf",
        "heavy_model": heavy,
    }
 
 
def run_weight_rank_tl(model, collapse_threshold: float = 0.1) -> dict:
    """Compute stable rank for Q/K/V/O and MLP matrices across all layers."""
    n_layers = int(getattr(model.cfg, "n_layers", 0) or len(model.blocks))
    matrices: list[dict] = []
 
    for layer in range(n_layers):
        block = model.blocks[layer]
        attn = block.attn
        mlp = block.mlp
        specs: list[tuple[str, torch.Tensor]] = [
            ("Q", attn.W_Q),
            ("K", attn.W_K),
            ("V", attn.W_V),
            ("O", attn.W_O),
            ("MLP_up", mlp.W_in),
            ("MLP_down", mlp.W_out),
        ]
        if hasattr(mlp, "W_gate"):
            specs.append(("MLP_gate", mlp.W_gate))
 
        for matrix_name, weight in specs:
            stable, nuclear = _matrix_ranks(weight.detach())
            matrices.append({
                "layer": layer,
                "matrix": matrix_name,
                "stable_rank": stable,
                "nuclear_norm_ratio": nuclear,
                "collapsed": stable < collapse_threshold,
            })
 
    collapsed_count = sum(1 for m in matrices if m["collapsed"])
    return {
        "n_layers": n_layers,
        "matrices": matrices,
        "collapse_threshold": collapse_threshold,
        "collapsed_count": collapsed_count,
        "backend": "transformer_lens",
    }
 
 
def run_weight_rank_analysis(model, collapse_threshold: float = 0.1) -> dict:
    """Compute stable rank for attention/MLP (TL) or HF-native weights."""
    if _is_hf_weight_model(model):
        return run_weight_rank_hf(model, collapse_threshold=collapse_threshold)
    return run_weight_rank_tl(model, collapse_threshold=collapse_threshold)