trajectory_analysis.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""Multi-checkpoint training trajectory: weight-diff summary per step vs base."""
 
from __future__ import annotations
 
import glob
from pathlib import Path
from typing import Any
 
from aquin.compute.sae_diff import load_checkpoint_state
from aquin.compute.weight_diff import run_weight_diff
 
 
def resolve_checkpoint_paths(*, pattern: str | None = None, directory: str | None = None) -> list[Path]:
    paths: list[Path] = []
    if pattern:
        for hit in glob.glob(str(Path(pattern).expanduser()), recursive=True):
            p = Path(hit)
            if p.is_file():
                paths.append(p)
    elif directory:
        root = Path(directory).expanduser()
        if not root.is_dir():
            raise FileNotFoundError(f"Directory not found: {root}")
        paths = [p for p in root.rglob("*.pt") if p.is_file()]
    else:
        raise ValueError("Provide --checkpoints <glob> or --dir <path>")
 
    if not paths:
        raise FileNotFoundError("No checkpoint files matched.")
 
    def _sort_key(p: Path) -> tuple[int, str]:
        try:
            _, step = load_checkpoint_state(p)
        except Exception:
            step = None
        return (step if step is not None else 10**9, p.name)
 
    return sorted(set(paths), key=_sort_key)
 
 
def run_trajectory_analysis(
    model_id: str,
    checkpoint_paths: list[str | Path],
    *,
    name_prefix: str | None = None,
) -> dict[str, Any]:
    steps: list[dict[str, Any]] = []
    prev_total: float | None = None
 
    for path in checkpoint_paths:
        ckpt = Path(path)
        label = f"{name_prefix}-{ckpt.stem}" if name_prefix else ckpt.stem
        try:
            weight = run_weight_diff(model_id, ckpt, checkpoint_name=label)
        except Exception as e:
            steps.append(
                {
                    "checkpoint": str(ckpt),
                    "name": label,
                    "error": str(e),
                }
            )
            continue
 
        if weight.get("error"):
            steps.append(
                {
                    "checkpoint": str(ckpt),
                    "name": label,
                    "error": weight.get("error"),
                }
            )
            continue
 
        total = float(weight.get("totalDeltaL2") or 0.0)
        delta_from_prev = round(total - prev_total, 6) if prev_total is not None else None
        prev_total = total
 
        steps.append(
            {
                "step": weight.get("trainingStep"),
                "checkpoint": str(ckpt),
                "name": label,
                "totalDeltaL2": weight.get("totalDeltaL2"),
                "maxDeltaL2": weight.get("maxDeltaL2"),
                "meanDeltaStableRank": weight.get("meanDeltaStableRank"),
                "nMatrices": weight.get("nMatrices"),
                "deltaFromPrevious": delta_from_prev,
                "layerProfile": weight.get("layerProfile") or [],
                "topChanged": (weight.get("topChanged") or [])[:5],
            }
        )
 
    valid = [s for s in steps if not s.get("error")]
    peak = max(valid, key=lambda s: float(s.get("totalDeltaL2") or 0), default=None)
 
    return {
        "schema_version": 1,
        "type": "trajectoryAnalysis",
        "baseModelId": model_id,
        "nCheckpoints": len(checkpoint_paths),
        "nAnalyzed": len(valid),
        "steps": steps,
        "peakStep": peak.get("step") if peak else None,
        "peakTotalDeltaL2": peak.get("totalDeltaL2") if peak else None,
    }
 
 
def run_trajectory_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"
    try:
        model_id = resolve_model_id(str(model_id))
    except ValueError as e:
        return {"error": str(e)}
 
    try:
        paths = resolve_checkpoint_paths(
            pattern=args.get("checkpoints"),
            directory=args.get("dir"),
        )
    except (FileNotFoundError, ValueError) as e:
        return {"error": str(e)}
 
    return run_trajectory_analysis(
        model_id,
        [str(p) for p in paths],
        name_prefix=args.get("name"),
    )