merge_analysis.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""Pre-merge LoRA analysis: weight-diff + rank/collapse signals + optional behavioral scores."""
 
from __future__ import annotations
 
from pathlib import Path
from typing import Any
 
from aquin.compute.weight_diff import run_weight_diff
 
 
def _collapse_signals(matrices: list[dict[str, Any]], mean_stable: float) -> list[dict[str, Any]]:
    threshold = max(3.0, mean_stable * 1.75)
    out: list[dict[str, Any]] = []
    for row in sorted(matrices, key=lambda m: m.get("delta_stable_rank", 0), reverse=True):
        rank = float(row.get("delta_stable_rank") or 0)
        if rank < threshold:
            continue
        out.append(
            {
                "layer": row.get("layer"),
                "matrix": row.get("matrix"),
                "module": row.get("module"),
                "deltaL2": row.get("delta_l2"),
                "deltaStableRank": rank,
                "deltaNuclearRatio": row.get("delta_nuclear_ratio"),
            }
        )
        if len(out) >= 12:
            break
    return out
 
 
def _merge_verdict(
    *,
    collapse_signals: list[dict[str, Any]],
    behavioral: dict[str, Any] | None,
    total_delta_l2: float,
) -> tuple[str, list[str]]:
    warnings: list[str] = []
    severity = 0
 
    if collapse_signals:
        n = len(collapse_signals)
        warnings.append(
            f"{n} matrix(ices) show elevated stable rank in ΔW "
            f"(possible rank expansion before merge)."
        )
        severity += 1 if n <= 3 else 2
 
    if total_delta_l2 > 50.0:
        warnings.append(f"Large aggregate ‖ΔW‖ ({total_delta_l2:.4g}) — verify adapter scale before merge.")
        severity += 1
 
    if behavioral:
        consistency = float(behavioral.get("consistencyScore") or 1.0)
        robustness = float(behavioral.get("robustnessScore") or 1.0)
        if consistency < 0.55:
            warnings.append(f"Low behavioral consistency ({consistency:.2f}) on probe generations.")
            severity += 2
        elif consistency < 0.7:
            warnings.append(f"Moderate behavioral drift ({consistency:.2f}) on probe generations.")
            severity += 1
        if robustness < 0.45:
            warnings.append(f"Low robustness score ({robustness:.2f}) across probe categories.")
            severity += 2
        elif robustness < 0.6:
            warnings.append(f"Moderate robustness score ({robustness:.2f}).")
            severity += 1
 
    if severity >= 3:
        return "fail", warnings
    if severity >= 1:
        return "warn", warnings
    return "pass", warnings
 
 
def run_merge_analysis(
    model_id: str,
    checkpoint_path: str | Path,
    *,
    checkpoint_name: str | None = None,
    prompts: list[str] | None = None,
    with_behavioral: bool = True,
) -> dict[str, Any]:
    ckpt = Path(checkpoint_path)
    name = checkpoint_name or ckpt.stem
 
    weight = run_weight_diff(model_id, ckpt, checkpoint_name=name)
    if weight.get("error"):
        return weight
 
    matrices = weight.get("matrices") or weight.get("topChanged") or []
    mean_stable = float(weight.get("meanDeltaStableRank") or 0.0)
    collapse = _collapse_signals(matrices, mean_stable)
 
    behavioral: dict[str, Any] | None = None
    behavioral_error: str | None = None
    if with_behavioral:
        try:
            from aquin.compute.model_diff import _run_model_diff
 
            behavioral = _run_model_diff(model_id, str(ckpt), prompts or [], n_prompts=5)
        except Exception as e:
            behavioral_error = str(e)
 
    verdict, warnings = _merge_verdict(
        collapse_signals=collapse,
        behavioral=behavioral,
        total_delta_l2=float(weight.get("totalDeltaL2") or 0.0),
    )
 
    return {
        "schema_version": 1,
        "type": "mergeAnalysis",
        "baseModelId": weight.get("baseModelId"),
        "ftCheckpointName": weight.get("ftCheckpointName"),
        "checkpointPath": weight.get("checkpointPath"),
        "modelMode": weight.get("modelMode"),
        "deltaMode": weight.get("deltaMode"),
        "trainingStep": weight.get("trainingStep"),
        "mergeVerdict": verdict,
        "warnings": warnings,
        "nMatrices": weight.get("nMatrices"),
        "totalDeltaL2": weight.get("totalDeltaL2"),
        "maxDeltaL2": weight.get("maxDeltaL2"),
        "meanDeltaStableRank": mean_stable,
        "collapseSignals": collapse,
        "topChanged": weight.get("topChanged") or [],
        "layerProfile": weight.get("layerProfile") or [],
        "behavioralDiff": behavioral,
        "behavioralError": behavioral_error,
        "withBehavioral": with_behavioral,
    }
 
 
def run_merge_analysis_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")
    if not checkpoint:
        return {"error": "Missing --checkpoint <path>. Run: aquin diff weight --help"}
 
    try:
        model_id = resolve_model_id(str(model_id))
    except ValueError as e:
        return {"error": str(e)}
 
    prompts_path = args.get("prompts")
    prompts = None
    if prompts_path:
        from aquin.compute.sae_diff import load_prompts
 
        prompts = load_prompts(str(prompts_path))
 
    return run_merge_analysis(
        model_id,
        checkpoint,
        checkpoint_name=args.get("name"),
        prompts=prompts,
        with_behavioral=not args.get("no_behavioral"),
    )