#!/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()