scripts/run_heapr.py

#!/usr/bin/env python3
"""Run HEAPr Stage A covariance and Stage B atomic scoring."""

from __future__ import annotations

import argparse
import shutil
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import numpy as np
from tqdm.auto import tqdm

from heapr.calibration import iter_token_batches, load_token_cache
from heapr.constants import DEFAULT_NUM_CHUNKS, DEFAULT_PRUNE_MODEL, DEFAULT_SEQ_LEN
from heapr.hf_cache import make_static_cache
from heapr.instrumentation import LagunaTraceContext
from heapr.model_utils import (
    build_max_memory,
    discover_sparse_layers,
    get_expert_tensors,
    get_model_layers,
    load_causal_lm,
    model_device_summary,
    validate_model_device_placement,
)
from heapr.scoring import (
    CovarianceStore,
    InMemoryCovarianceStore,
    accumulate_covariance_from_trace,
    build_global_score_artifacts,
    compute_atomic_scores_for_expert,
    compute_down_covariance_quadratic,
)
from heapr.utils import collect_hardware_metadata, require_torch, write_json


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-id", default=DEFAULT_PRUNE_MODEL)
    parser.add_argument("--revision")
    parser.add_argument("--cache-path", required=True)
    parser.add_argument("--output-dir", required=True)
    parser.add_argument("--stage", choices=["covariance", "scores", "all"], default="all")
    parser.add_argument("--layer-mode", choices=["sequential", "all-at-once"], default="sequential")
    parser.add_argument("--batch-size", type=int, default=1)
    parser.add_argument("--max-chunks", type=int, default=DEFAULT_NUM_CHUNKS)
    parser.add_argument("--seq-len", type=int, default=DEFAULT_SEQ_LEN)
    parser.add_argument("--dtype", default="bfloat16")
    parser.add_argument("--attn-implementation")
    parser.add_argument("--gpu-memory-per-device")
    parser.add_argument("--max-gpu-memory")
    parser.add_argument("--max-cpu-memory")
    parser.add_argument("--offload-folder")
    parser.add_argument("--allow-cpu-offload", action="store_true")
    parser.add_argument("--covariance-dir")
    parser.add_argument("--keep-covariance", action="store_true")
    parser.add_argument("--atom-chunk", type=int, default=32)
    parser.add_argument("--loss-token-chunk", type=int, default=256)
    parser.add_argument("--max-sparse-layers", type=int)
    parser.add_argument("--layer-window-size", type=int, default=8)
    parser.add_argument("--covariance-accumulation", choices=["cpu", "device"], default="device")
    parser.add_argument("--no-fused-scoring", action="store_true")
    parser.add_argument("--random-seed", type=int, default=0)
    parser.add_argument("--no-gradient-checkpointing", action="store_true")
    parser.add_argument("--cache-implementation", default="static")
    parser.add_argument("--no-cache", action="store_true")
    return parser.parse_args()


def _first_device(model):
    return next(model.parameters()).device


def compute_chunked_lm_loss(
    model,
    input_ids,
    *,
    token_chunk: int,
    use_cache: bool = True,
    cache_implementation: str | None = "static",
):
    torch = require_torch()
    base_model = getattr(model, "model", None)
    lm_head = getattr(model, "lm_head", None)
    if base_model is None or lm_head is None:
        kwargs = {"input_ids": input_ids, "labels": input_ids, "use_cache": use_cache}
        if use_cache and cache_implementation == "static":
            kwargs["past_key_values"] = make_static_cache(model, max_cache_len=int(input_ids.shape[1]))
        return model(**kwargs).loss

    kwargs = {"input_ids": input_ids, "use_cache": use_cache, "return_dict": True}
    if use_cache and cache_implementation == "static":
        kwargs["past_key_values"] = make_static_cache(model, max_cache_len=int(input_ids.shape[1]))
    outputs = base_model(**kwargs)
    hidden_states = outputs.last_hidden_state
    shift_hidden = hidden_states[:, :-1, :]
    shift_labels = input_ids[:, 1:]
    flat_hidden = shift_hidden.reshape(-1, shift_hidden.shape[-1])
    flat_labels = shift_labels.reshape(-1)
    if token_chunk <= 0 or token_chunk >= flat_labels.numel():
        logits = lm_head(flat_hidden).float()
        labels = flat_labels.to(logits.device)
        return torch.nn.functional.cross_entropy(logits, labels)
    total_loss = None
    total_tokens = flat_labels.numel()
    for start in range(0, total_tokens, token_chunk):
        end = min(start + token_chunk, total_tokens)
        logits = lm_head(flat_hidden[start:end]).float()
        labels = flat_labels[start:end].to(logits.device)
        loss = torch.nn.functional.cross_entropy(logits, labels, reduction="sum")
        total_loss = loss if total_loss is None else total_loss + loss
    return total_loss / total_tokens


def _remove_layer_covariances(cov_dir: Path, sparse_idx: int) -> None:
    for path in cov_dir.glob(f"cov_l{sparse_idx:02d}_e*.mmap"):
        path.unlink()


def _covariance_dir(output_dir: Path, args) -> Path:
    return Path(args.covariance_dir) if args.covariance_dir else output_dir / "covariance"


def _sparse_layers_for_run(model, args):
    sparse_layers = discover_sparse_layers(model)
    if args.max_sparse_layers:
        sparse_layers = sparse_layers[: args.max_sparse_layers]
    return sparse_layers


def _layer_windows(sparse_layers, window_size: int):
    window_size = max(int(window_size), 1)
    for start in range(0, len(sparse_layers), window_size):
        yield sparse_layers[start : start + window_size]


def _token_chunks_for_run(token_chunks, args):
    chunks = token_chunks[: args.max_chunks] if args.max_chunks else token_chunks
    if args.seq_len and args.seq_len > 0:
        chunks = chunks[:, : args.seq_len]
    return chunks


def _score_array_shape(sparse_layers):
    first = sparse_layers[0]
    return len(sparse_layers), first.num_experts, first.routed_width


def accumulate_activation_stats_from_trace(trace, layers, activation_sums, token_counts, *, moe_scale: float) -> None:
    torch = require_torch()
    with torch.no_grad():
        for record in trace.routes:
            if record.hidden_states is None:
                raise RuntimeError("trace did not store hidden states for fused scoring")
            mlp = getattr(layers[record.layer_model_idx], "mlp")
            gate_up_all, _ = get_expert_tensors(mlp)
            selected = record.selected_experts
            routing = record.routing_weights
            for expert in torch.unique(selected).detach().cpu().tolist():
                matches = selected == int(expert)
                token_rows = matches.any(dim=-1)
                n_tokens = int(token_rows.sum().detach().cpu())
                if n_tokens == 0:
                    continue
                route_weights = (routing * matches.to(routing.dtype)).sum(dim=-1)[token_rows]
                gate_up = gate_up_all[int(expert)]
                hidden = record.hidden_states[token_rows].to(device=gate_up.device, dtype=gate_up.dtype)
                route = route_weights.to(device=gate_up.device, dtype=gate_up.dtype)
                routed_width = gate_up.shape[0] // 2
                gate = gate_up[:routed_width, :]
                up = gate_up[routed_width:, :]
                gate_values = torch.nn.functional.silu(hidden @ gate.T)
                up_values = hidden @ up.T
                activation_scale = (gate_values * up_values).float() * route[:, None].float() * float(moe_scale)
                activation_sums[record.layer_sparse_idx, int(expert)] += (
                    activation_scale.square().sum(dim=0).detach().cpu().numpy().astype(np.float64)
                )
                token_counts[record.layer_sparse_idx, int(expert)] += n_tokens


def finalize_scores_from_activation_stats(
    model,
    cov_store,
    activation_sums: np.ndarray,
    token_counts: np.ndarray,
    args,
) -> np.ndarray:
    sparse_layers = discover_sparse_layers(model)
    layers = get_model_layers(model)
    atomic_scores = np.zeros_like(activation_sums, dtype=np.float32)
    for info in sparse_layers[: activation_sums.shape[0]]:
        mlp = getattr(layers[info.model_layer_idx], "mlp")
        _, down_all = get_expert_tensors(mlp)
        for expert in range(info.num_experts):
            count = int(token_counts[info.sparse_idx, expert])
            if count == 0 or cov_store.counts.get((info.sparse_idx, expert), 0) == 0:
                continue
            cov = cov_store.normalized(info.sparse_idx, expert)
            cov_quadratic = compute_down_covariance_quadratic(
                down_all[expert],
                cov,
                atom_chunk=args.atom_chunk,
            )
            activation_mean = activation_sums[info.sparse_idx, expert] / count
            atomic_scores[info.sparse_idx, expert] = (
                0.5 * activation_mean * cov_quadratic.detach().cpu().numpy().astype(np.float64)
            ).astype(np.float32)
    return atomic_scores


def run_covariance(
    model,
    token_chunks,
    output_dir: Path,
    args,
    *,
    target_sparse_idx: int | None = None,
    target_sparse_indices: set[int] | None = None,
    store=None,
    store_hidden: bool = False,
    forward_callback=None,
) -> None:
    torch = require_torch()
    cov_dir = _covariance_dir(output_dir, args)
    sparse_layers = discover_sparse_layers(model)
    hidden_size = sparse_layers[0].hidden_size
    store = store if store is not None else CovarianceStore(cov_dir, hidden_size=hidden_size)
    chunks = _token_chunks_for_run(token_chunks, args)
    device = _first_device(model)
    original_requires_grad = [parameter.requires_grad for parameter in model.parameters()]
    for parameter in model.parameters():
        parameter.requires_grad_(False)
    if target_sparse_indices is not None:
        layer_filter = set(target_sparse_indices)
    elif target_sparse_idx is not None:
        layer_filter = {target_sparse_idx}
    else:
        layer_filter = None
    force_single_layer_leaf = layer_filter is not None and len(layer_filter) == 1
    force_first_layer_leaf = not force_single_layer_leaf
    layer_desc = (
        ",".join(str(idx) for idx in sorted(layer_filter))
        if layer_filter is not None
        else "all"
    )
    try:
        with LagunaTraceContext(
            model,
            store_hidden=store_hidden,
            sparse_layer_filter=layer_filter,
            force_output_requires_grad=force_single_layer_leaf,
            force_first_output_requires_grad=force_first_layer_leaf,
        ) as trace:
            num_batches = int(np.ceil(chunks.shape[0] / args.batch_size))
            for batch_idx, batch in enumerate(
                tqdm(
                    iter_token_batches(chunks, batch_size=args.batch_size),
                    total=num_batches,
                    desc=f"cov layer {layer_desc}",
                )
            ):
                trace.clear()
                model.zero_grad(set_to_none=True)
                input_ids = torch.as_tensor(np.array(batch, copy=True), device=device)
                batch_start = time.perf_counter()
                print(f"[covariance] batch={batch_idx + 1}/{num_batches} forward+loss start", flush=True)
                loss = compute_chunked_lm_loss(
                    model,
                    input_ids,
                    token_chunk=args.loss_token_chunk,
                    use_cache=not args.no_cache,
                    cache_implementation=args.cache_implementation,
                )
                if forward_callback is not None:
                    forward_callback(trace)
                forward_elapsed = time.perf_counter() - batch_start
                print(
                    f"[covariance] batch={batch_idx + 1}/{num_batches} backward start "
                    f"loss={float(loss.detach().cpu()):.6f} forward_loss_s={forward_elapsed:.1f}",
                    flush=True,
                )
                backward_start = time.perf_counter()
                loss.backward()
                backward_elapsed = time.perf_counter() - backward_start
                print(
                    f"[covariance] batch={batch_idx + 1}/{num_batches} accumulate start "
                    f"routes={len(trace.routes)} backward_s={backward_elapsed:.1f}",
                    flush=True,
                )
                accumulate_start = time.perf_counter()
                accumulate_covariance_from_trace(
                    trace,
                    store,
                    show_progress=len(trace.routes) > 1,
                    device_accumulation=args.covariance_accumulation == "device",
                )
                accumulate_elapsed = time.perf_counter() - accumulate_start
                total_elapsed = time.perf_counter() - batch_start
                print(
                    f"[covariance] batch={batch_idx + 1}/{num_batches} done "
                    f"accumulate_s={accumulate_elapsed:.1f} total_s={total_elapsed:.1f}",
                    flush=True,
                )
    finally:
        for parameter, requires_grad in zip(model.parameters(), original_requires_grad):
            parameter.requires_grad_(requires_grad)
    store.save_counts()
    write_json(
        cov_dir / "metadata.json",
        {
            "num_sparse_layers": len(sparse_layers),
            "hidden_size": hidden_size,
            "num_chunks": int(chunks.shape[0]),
            "seq_len": int(chunks.shape[1]),
            "batch_size": args.batch_size,
            "device_summary": model_device_summary(model),
            "target_sparse_idx": target_sparse_idx,
            "target_sparse_indices": sorted(layer_filter) if layer_filter is not None else None,
            "gradient_mode": (
                "single_layer_detached_leaf"
                if force_single_layer_leaf
                else "first_traced_sparse_detached_leaf"
            ),
        },
    )


def run_covariance_sequential(model, token_chunks, output_dir: Path, args) -> None:
    cov_dir = _covariance_dir(output_dir, args)
    cov_dir.mkdir(parents=True, exist_ok=True)
    sparse_layers = _sparse_layers_for_run(model, args)
    chunks = _token_chunks_for_run(token_chunks, args)
    combined_counts: dict[tuple[int, int], int] = {}
    for info in sparse_layers:
        print(f"[covariance sequential] sparse_layer={info.sparse_idx} model_layer={info.model_layer_idx}")
        _remove_layer_covariances(cov_dir, info.sparse_idx)
        run_covariance(model, token_chunks, output_dir, args, target_sparse_idx=info.sparse_idx)
        layer_store = CovarianceStore(cov_dir, hidden_size=info.hidden_size)
        layer_store.load_counts()
        combined_counts.update(layer_store.counts)
    write_json(
        cov_dir / "counts.json",
        {
            f"{layer}:{expert}": count
            for (layer, expert), count in sorted(combined_counts.items())
        },
    )
    write_json(
        cov_dir / "metadata.json",
        {
            "num_sparse_layers": len(sparse_layers),
            "hidden_size": sparse_layers[0].hidden_size,
            "num_chunks": int(chunks.shape[0]),
            "seq_len": int(chunks.shape[1]),
            "batch_size": args.batch_size,
            "device_summary": model_device_summary(model),
            "target_sparse_idx": None,
            "gradient_mode": "sequential_single_layer_detached_leaf",
        },
    )


def run_scores(
    model,
    token_chunks,
    output_dir: Path,
    args,
    *,
    target_sparse_idx: int | None = None,
    target_sparse_indices: set[int] | None = None,
    build_artifacts: bool = True,
    cov_store=None,
) -> tuple[np.ndarray, np.ndarray]:
    torch = require_torch()
    sparse_layers = discover_sparse_layers(model)
    layers = get_model_layers(model)
    hidden_size = sparse_layers[0].hidden_size
    routed_width = sparse_layers[0].routed_width
    num_experts = sparse_layers[0].num_experts
    cov_store = cov_store if cov_store is not None else CovarianceStore(_covariance_dir(output_dir, args), hidden_size=hidden_size)
    cov_store.load_counts()

    score_sums = np.zeros((len(sparse_layers), num_experts, routed_width), dtype=np.float64)
    token_counts = np.zeros((len(sparse_layers), num_experts), dtype=np.int64)
    chunks = _token_chunks_for_run(token_chunks, args)
    device = _first_device(model)
    moe_scale = float(getattr(getattr(model, "config", None), "moe_routed_scaling_factor", 1.0))

    if target_sparse_indices is not None:
        layer_filter = set(target_sparse_indices)
    elif target_sparse_idx is not None:
        layer_filter = {target_sparse_idx}
    else:
        layer_filter = None
    layer_desc = (
        ",".join(str(idx) for idx in sorted(layer_filter))
        if layer_filter is not None
        else "all"
    )
    with LagunaTraceContext(model, store_hidden=True, sparse_layer_filter=layer_filter) as trace:
        num_batches = int(np.ceil(chunks.shape[0] / args.batch_size))
        for batch_idx, batch in enumerate(
            tqdm(
                iter_token_batches(chunks, batch_size=args.batch_size),
                total=num_batches,
                desc=f"score layer {layer_desc}",
            )
        ):
            trace.clear()
            input_ids = torch.as_tensor(np.array(batch, copy=True), device=device)
            with torch.no_grad():
                forward_kwargs = {"input_ids": input_ids, "use_cache": not args.no_cache}
                if not args.no_cache and args.cache_implementation == "static":
                    forward_kwargs["past_key_values"] = make_static_cache(
                        model,
                        max_cache_len=int(input_ids.shape[1]),
                    )
                model(**forward_kwargs)
            for record in trace.routes:
                mlp = getattr(layers[record.layer_model_idx], "mlp")
                gate_up_all, down_all = get_expert_tensors(mlp)
                selected = record.selected_experts
                routing = record.routing_weights
                if record.hidden_states is None:
                    raise RuntimeError("trace did not store hidden states for scoring")
                for expert in torch.unique(selected).detach().cpu().tolist():
                    matches = selected == int(expert)
                    token_rows = matches.any(dim=-1)
                    n_tokens = int(token_rows.sum().detach().cpu())
                    if n_tokens == 0 or cov_store.counts.get((record.layer_sparse_idx, int(expert)), 0) == 0:
                        continue
                    route_weights = (routing * matches.to(routing.dtype)).sum(dim=-1)[token_rows]
                    cov = cov_store.normalized(record.layer_sparse_idx, int(expert))
                    expert_scores = compute_atomic_scores_for_expert(
                        record.hidden_states[token_rows],
                        route_weights,
                        gate_up_all[int(expert)],
                        down_all[int(expert)],
                        cov,
                        moe_scale=moe_scale,
                        atom_chunk=args.atom_chunk,
                    )
                    score_sums[record.layer_sparse_idx, int(expert)] += (
                        expert_scores.numpy().astype(np.float64) * n_tokens
                    )
                    token_counts[record.layer_sparse_idx, int(expert)] += n_tokens
            if (batch_idx + 1) % 10 == 0:
                print(f"[scores] batch={batch_idx + 1}")

    atomic_scores = score_sums / np.maximum(token_counts[:, :, None], 1)
    if build_artifacts:
        np.save(output_dir / "atomic_token_counts.npy", token_counts)
        build_global_score_artifacts(
            atomic_scores.astype(np.float32),
            output_dir,
            random_seed=args.random_seed,
        )
    return atomic_scores, token_counts


def run_all_sequential(model, token_chunks, output_dir: Path, args) -> None:
    sparse_layers = _sparse_layers_for_run(model, args)
    num_layers, num_experts, routed_width = _score_array_shape(sparse_layers)
    atomic_scores = np.zeros((num_layers, num_experts, routed_width), dtype=np.float32)
    token_counts = np.zeros((num_layers, num_experts), dtype=np.int64)
    cov_dir = _covariance_dir(output_dir, args)
    cov_dir.mkdir(parents=True, exist_ok=True)
    moe_scale = float(getattr(getattr(model, "config", None), "moe_routed_scaling_factor", 1.0))
    layers = get_model_layers(model)

    for window in _layer_windows(sparse_layers, args.layer_window_size):
        sparse_indices = {info.sparse_idx for info in window}
        desc = ",".join(str(info.sparse_idx) for info in window)
        model_desc = ",".join(str(info.model_layer_idx) for info in window)
        print(f"[sequential] sparse_layers={desc} model_layers={model_desc}")
        for info in window:
            _remove_layer_covariances(cov_dir, info.sparse_idx)
        layer_cov_store = InMemoryCovarianceStore(hidden_size=window[0].hidden_size)
        if args.no_fused_scoring:
            run_covariance(
                model,
                token_chunks,
                output_dir,
                args,
                target_sparse_indices=sparse_indices,
                store=layer_cov_store,
            )
            layer_scores, layer_counts = run_scores(
                model,
                token_chunks,
                output_dir,
                args,
                target_sparse_indices=sparse_indices,
                build_artifacts=False,
                cov_store=layer_cov_store,
            )
        else:
            activation_sums = np.zeros_like(atomic_scores, dtype=np.float64)
            layer_counts = np.zeros_like(token_counts)

            def collect_activation_stats(trace) -> None:
                accumulate_activation_stats_from_trace(
                    trace,
                    layers,
                    activation_sums,
                    layer_counts,
                    moe_scale=moe_scale,
                )

            run_covariance(
                model,
                token_chunks,
                output_dir,
                args,
                target_sparse_indices=sparse_indices,
                store=layer_cov_store,
                store_hidden=True,
                forward_callback=collect_activation_stats,
            )
            layer_scores = finalize_scores_from_activation_stats(
                model,
                layer_cov_store,
                activation_sums,
                layer_counts,
                args,
            )
        for info in window:
            atomic_scores[info.sparse_idx] = layer_scores[info.sparse_idx].astype(np.float32)
            token_counts[info.sparse_idx] = layer_counts[info.sparse_idx]
            _remove_layer_covariances(cov_dir, info.sparse_idx)

    np.save(output_dir / "atomic_token_counts.npy", token_counts)
    build_global_score_artifacts(atomic_scores, output_dir, random_seed=args.random_seed)
    shutil.rmtree(cov_dir, ignore_errors=True)


def run_all_in_memory(model, token_chunks, output_dir: Path, args) -> None:
    sparse_layers = _sparse_layers_for_run(model, args)
    hidden_size = sparse_layers[0].hidden_size
    cov_store = InMemoryCovarianceStore(hidden_size=hidden_size)
    sparse_indices = {info.sparse_idx for info in sparse_layers}
    run_covariance(
        model,
        token_chunks,
        output_dir,
        args,
        target_sparse_indices=sparse_indices,
        store=cov_store,
    )
    run_scores(model, token_chunks, output_dir, args, target_sparse_indices=sparse_indices, cov_store=cov_store)
    shutil.rmtree(_covariance_dir(output_dir, args), ignore_errors=True)


def main() -> None:
    args = parse_args()
    if args.offload_folder and not args.allow_cpu_offload:
        raise ValueError("--offload-folder requires --allow-cpu-offload")
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    write_json(output_dir / "hardware.json", collect_hardware_metadata())
    token_chunks = load_token_cache(args.cache_path)
    max_memory = build_max_memory(
        gpu_memory_per_device=args.gpu_memory_per_device,
        max_gpu_memory=args.max_gpu_memory,
        max_cpu_memory=args.max_cpu_memory,
        allow_cpu_offload=args.allow_cpu_offload,
    )
    torch = require_torch()
    requested_gpu_count = torch.cuda.device_count() if args.gpu_memory_per_device else None
    model = load_causal_lm(
        args.model_id,
        revision=args.revision,
        dtype=args.dtype,
        max_memory=max_memory,
        offload_folder=args.offload_folder if args.allow_cpu_offload else None,
        attn_implementation=args.attn_implementation,
        use_cache=not args.no_cache,
        cache_implementation=args.cache_implementation,
        output_router_logits=False,
    )
    validate_model_device_placement(
        model,
        allow_cpu_offload=args.allow_cpu_offload,
        requested_gpu_count=requested_gpu_count,
    )
    if not args.no_gradient_checkpointing and hasattr(model, "gradient_checkpointing_enable"):
        model.gradient_checkpointing_enable()
    write_json(output_dir / "model_device_summary.json", model_device_summary(model))
    if args.stage == "all" and args.layer_mode == "sequential":
        run_all_sequential(model, token_chunks, output_dir, args)
        return
    if args.stage == "all" and args.layer_mode == "all-at-once":
        run_all_in_memory(model, token_chunks, output_dir, args)
        return
    if args.stage == "covariance" and args.layer_mode == "sequential":
        run_covariance_sequential(model, token_chunks, output_dir, args)
        return
    if args.stage in {"covariance", "all"}:
        run_covariance(model, token_chunks, output_dir, args)
    if args.stage in {"scores", "all"}:
        run_scores(model, token_chunks, output_dir, args)
    if args.stage == "all" and not args.keep_covariance:
        shutil.rmtree(_covariance_dir(output_dir, args), ignore_errors=True)


if __name__ == "__main__":
    main()