weight_diff.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""Weight delta analysis: base vs fine-tuned checkpoint (LLM)."""
 
from __future__ import annotations
 
import re
from collections import defaultdict
from pathlib import Path
from typing import Any
 
import torch
 
from aquin.compute.sae_diff import load_checkpoint_state
from aquin.compute.weight_rank import _matrix_ranks, _to_2d
 
 
_LAYER_RE = re.compile(r"(?:^|\.)(?:layers?|layer|block)\.?(\d+)(?:\.|$)", re.I)
_DEFAULT_LORA_R = 8
_DEFAULT_LORA_ALPHA = 16.0
_SMOKE_LORA_PATH = Path.home() / ".aquin" / "smoke" / "weight-diff-lora.pt"
 
 
def _layer_from_key(key: str) -> int:
    m = _LAYER_RE.search(key.replace("_", "."))
    return int(m.group(1)) if m else -1
 
 
def _matrix_label_from_key(key: str) -> str:
    kl = key.lower()
    for tag in ("q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", "query", "key", "value"):
        if tag in kl:
            return tag
    if "mlp" in kl and "gate" in kl:
        return "mlp_gate"
    if "mlp" in kl and "down" in kl:
        return "mlp_down"
    if "mlp" in kl:
        return "mlp"
    if "attn" in kl:
        return "attn"
    parts = key.split(".")
    return parts[-2] if len(parts) >= 2 else parts[-1]
 
 
def ensure_smoke_lora_checkpoint() -> Path:
    """Tiny LoRA state_dict so desktop smoke can run without a real fine-tune."""
    dest = _SMOKE_LORA_PATH
    if dest.is_file():
        return dest
    dest.parent.mkdir(parents=True, exist_ok=True)
    r, d = _DEFAULT_LORA_R, 32
    state = {
        "model.layers.0.self_attn.q_proj.lora_A": torch.randn(r, d) * 0.05,
        "model.layers.0.self_attn.q_proj.lora_B": torch.randn(d, r) * 0.05,
        "model.layers.1.self_attn.v_proj.lora_A": torch.randn(r, d) * 0.05,
        "model.layers.1.self_attn.v_proj.lora_B": torch.randn(d, r) * 0.05,
    }
    torch.save(state, dest)
    return dest
 
 
def _has_lora_keys(state: dict[str, Any]) -> bool:
    return any("lora_a" in k.lower() or "lora_b" in k.lower() for k in state)
 
 
def _matrix_entry(
    *,
    layer: int,
    matrix: str,
    module: str,
    delta: torch.Tensor,
    base_norm: float | None = None,
) -> dict[str, Any]:
    l2 = float(delta.norm().item())
    rel = (l2 / max(base_norm, 1e-9)) if base_norm is not None and base_norm > 1e-9 else 0.0
    stable, nuclear = _matrix_ranks(delta)
    return {
        "layer": layer,
        "matrix": matrix,
        "module": module,
        "delta_l2": round(l2, 6),
        "delta_l2_relative": round(rel, 6),
        "delta_stable_rank": stable,
        "delta_nuclear_ratio": nuclear,
    }
 
 
def _layer_profile(matrices: list[dict[str, Any]]) -> list[dict[str, Any]]:
    totals: dict[int, float] = defaultdict(float)
    counts: dict[int, int] = defaultdict(int)
    for row in matrices:
        layer = int(row["layer"])
        if layer < 0:
            continue
        totals[layer] += float(row["delta_l2"])
        counts[layer] += 1
    return [
        {
            "layer": layer,
            "delta_l2": round(totals[layer], 6),
            "n_matrices": counts[layer],
        }
        for layer in sorted(totals)
    ]
 
 
def _finalize_payload(
    *,
    model_id: str,
    mode: str,
    checkpoint_name: str,
    checkpoint_path: str,
    delta_mode: str,
    matrices: list[dict[str, Any]],
    step: int | None = None,
) -> dict[str, Any]:
    sorted_m = sorted(matrices, key=lambda x: x["delta_l2"], reverse=True)
    total_l2 = round(sum(m["delta_l2"] for m in matrices), 6)
    max_l2 = round(max((m["delta_l2"] for m in matrices), default=0.0), 6)
    mean_stable = (
        round(sum(m["delta_stable_rank"] for m in matrices) / len(matrices), 4)
        if matrices else 0.0
    )
    return {
        "schema_version": 1,
        "type": "weightDiff",
        "baseModelId": model_id,
        "ftCheckpointName": checkpoint_name,
        "checkpointPath": checkpoint_path,
        "modelMode": mode,
        "deltaMode": delta_mode,
        "trainingStep": step,
        "nMatrices": len(matrices),
        "totalDeltaL2": total_l2,
        "maxDeltaL2": max_l2,
        "meanDeltaStableRank": mean_stable,
        "layerProfile": _layer_profile(matrices),
        "topChanged": sorted_m[:25],
        "matrices": sorted_m,
    }
 
 
def _iter_tl_matrix_pairs(tl_base: Any, tl_ft: Any) -> list[tuple[int, str, str, torch.Tensor, torch.Tensor]]:
    n_layers = int(getattr(tl_base.cfg, "n_layers", 0) or len(tl_base.blocks))
    pairs: list[tuple[int, str, str, torch.Tensor, torch.Tensor]] = []
    for layer in range(n_layers):
        block_b = tl_base.blocks[layer]
        block_f = tl_ft.blocks[layer]
        attn = block_b.attn
        attn_f = block_f.attn
        mlp_b = block_b.mlp
        mlp_f = block_f.mlp
        specs: list[tuple[str, torch.Tensor, torch.Tensor]] = [
            ("Q", attn.W_Q, attn_f.W_Q),
            ("K", attn.W_K, attn_f.W_K),
            ("V", attn.W_V, attn_f.W_V),
            ("O", attn.W_O, attn_f.W_O),
            ("MLP_up", mlp_b.W_in, mlp_f.W_in),
            ("MLP_down", mlp_b.W_out, mlp_f.W_out),
        ]
        if hasattr(mlp_b, "W_gate"):
            specs.append(("MLP_gate", mlp_b.W_gate, mlp_f.W_gate))
        for name, wb, wf in specs:
            pairs.append((layer, name, f"blocks.{layer}.{name}", wb.detach(), wf.detach()))
    return pairs
 
 
def run_llm_hf_weight_diff(
    model_id: str,
    checkpoint_path: str | Path,
    *,
    checkpoint_name: str = "checkpoint",
) -> dict[str, Any]:
    from transformers import AutoModelForCausalLM
 
    from aquin.compute.model_loader import get_config, resolve_model_id
    from aquin.compute.sae_diff import load_checkpoint_state
 
    short = resolve_model_id(model_id)
    cfg = get_config(short)
    state, step = load_checkpoint_state(checkpoint_path)
    if _has_lora_keys(state):
        return run_llm_lora_weight_diff(
            short,
            state,
            checkpoint_name=checkpoint_name,
            checkpoint_path=str(checkpoint_path),
            step=step,
        )
 
    hf_name = cfg["hf_name"]
    trust = bool(cfg.get("trust_remote_code", False))
    base_model = AutoModelForCausalLM.from_pretrained(
        hf_name,
        torch_dtype=torch.float32,
        device_map="cpu",
        low_cpu_mem_usage=True,
        trust_remote_code=trust,
    )
    base_sd = {k: v.detach().float().cpu() for k, v in base_model.state_dict().items()}
    del base_model
 
    matrices: list[dict[str, Any]] = []
    for key, ft_w in state.items():
        base_w = base_sd.get(key)
        if base_w is None or not isinstance(ft_w, torch.Tensor):
            continue
        wf = ft_w.detach().float().cpu()
        if base_w.shape != wf.shape or base_w.ndim < 2:
            continue
        delta = wf - base_w
        layer = _layer_from_key(key)
        matrix = _matrix_label_from_key(key)
        base_norm = float(_to_2d(base_w).norm().item())
        matrices.append(
            _matrix_entry(
                layer=layer,
                matrix=matrix,
                module=key,
                delta=delta,
                base_norm=base_norm,
            )
        )
    return _finalize_payload(
        model_id=short,
        mode="llm",
        checkpoint_name=checkpoint_name,
        checkpoint_path=str(checkpoint_path),
        delta_mode="hf_state_dict",
        matrices=matrices,
        step=step,
    )
 
 
def run_llm_lora_weight_diff(
    model_id: str,
    state_dict: dict[str, Any],
    *,
    checkpoint_name: str = "checkpoint",
    checkpoint_path: str = "",
    step: int | None = None,
    lora_r: int = _DEFAULT_LORA_R,
    lora_alpha: float = _DEFAULT_LORA_ALPHA,
) -> dict[str, Any]:
    pairs: dict[str, dict[str, torch.Tensor]] = {}
    for key, tensor in state_dict.items():
        kl = key.lower()
        if ".lora_a" in kl:
            prefix = key.split(".lora_A")[0].split(".lora_a")[0]
            pairs.setdefault(prefix, {})["A"] = tensor.float()
        elif ".lora_b" in kl:
            prefix = key.split(".lora_B")[0].split(".lora_b")[0]
            pairs.setdefault(prefix, {})["B"] = tensor.float()
 
    scale = lora_alpha / max(lora_r, 1)
    matrices: list[dict[str, Any]] = []
    for prefix, ab in pairs.items():
        a = ab.get("A")
        b = ab.get("B")
        if a is None or b is None:
            continue
        if b.ndim == 2 and a.ndim == 2:
            delta = (b @ a) * scale
        else:
            delta = (b.reshape(-1, 1) @ a.reshape(1, -1)) * scale
        layer = _layer_from_key(prefix)
        matrix = _matrix_label_from_key(prefix)
        matrices.append(
            _matrix_entry(
                layer=layer,
                matrix=matrix,
                module=prefix,
                delta=delta.cpu(),
                base_norm=None,
            )
        )
 
    return _finalize_payload(
        model_id=model_id,
        mode="llm",
        checkpoint_name=checkpoint_name,
        checkpoint_path=checkpoint_path,
        delta_mode="lora_effective",
        matrices=matrices,
        step=step,
    )
 
 
def run_weight_diff(
    model_id: str,
    checkpoint_path: str | Path,
    *,
    checkpoint_name: str | None = None,
) -> dict[str, Any]:
    from aquin.compute.model_loader import resolve_model_id
 
    ckpt = Path(checkpoint_path)
    if not ckpt.exists():
        raise FileNotFoundError(f"Checkpoint not found: {ckpt}")
 
    name = checkpoint_name or ckpt.stem
    short = resolve_model_id(model_id)
    return run_llm_hf_weight_diff(short, ckpt, checkpoint_name=name)
 
 
def run_weight_diff_from_args(args: dict[str, Any]) -> dict[str, Any]:
    from aquin.compute.model_loader import get_active_model_id, resolve_model_id
 
    model_id = args.get("model_id") or get_active_model_id() or "llama-3.2-1b"
    checkpoint = args.get("checkpoint")
    raw = str(checkpoint or "").strip()
    if not raw or raw.lower() in ("none", "null", "default", "-"):
        checkpoint = str(ensure_smoke_lora_checkpoint())
 
    try:
        model_id = resolve_model_id(str(model_id))
    except ValueError as exc:
        return {"error": str(exc)}
 
    try:
        return run_weight_diff(
            model_id,
            checkpoint,
            checkpoint_name=args.get("name"),
        )
    except Exception as exc:
        return {"error": str(exc)}