| |
| """Evaluate held-out loss after applying HEAPr pruning masks in memory.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import sys |
| from pathlib import Path |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) |
|
|
| import numpy as np |
|
|
| from heapr.constants import DEFAULT_GROUP_WIDTH |
| from heapr.eval.loss import evaluate_token_cache |
| from heapr.model_utils import ( |
| build_max_memory, |
| load_causal_lm, |
| validate_model_device_placement, |
| ) |
| from heapr.prune import ( |
| apply_atomic_mask_to_model, |
| apply_group_mask_to_model, |
| atomic_mask_from_scores, |
| group_mask_from_scores, |
| ) |
| from heapr.utils import require_torch, write_json |
|
|
|
|
| def parse_args(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--model-id", required=True) |
| parser.add_argument("--cache-path", required=True) |
| parser.add_argument("--scores-path", required=True) |
| parser.add_argument("--output-dir", required=True) |
| parser.add_argument("--mode", choices=["atomic", "group"], default="atomic") |
| parser.add_argument("--group-indices") |
| parser.add_argument("--group-width", type=int, default=DEFAULT_GROUP_WIDTH) |
| parser.add_argument("--ratios", nargs="+", type=float, default=[0.10, 0.20, 0.40]) |
| parser.add_argument("--revision") |
| parser.add_argument("--batch-size", type=int, default=1) |
| parser.add_argument("--max-chunks", type=int) |
| parser.add_argument("--dtype", default="bfloat16") |
| 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("--cache-implementation", default="static") |
| parser.add_argument("--no-cache", action="store_true") |
| return parser.parse_args() |
|
|
|
|
| def main() -> None: |
| args = parse_args() |
| if args.offload_folder and not args.allow_cpu_offload: |
| raise ValueError("--offload-folder requires --allow-cpu-offload") |
| if args.mode == "group" and not args.group_indices: |
| raise ValueError("--group-indices is required for grouped pruning") |
|
|
| scores = np.load(args.scores_path) |
| group_indices = np.load(args.group_indices) if args.group_indices else None |
| output_dir = Path(args.output_dir) |
| output_dir.mkdir(parents=True, exist_ok=True) |
|
|
| 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 |
|
|
| for ratio in args.ratios: |
| if torch.cuda.is_available(): |
| torch.cuda.empty_cache() |
| for device_idx in range(torch.cuda.device_count()): |
| try: |
| torch.cuda.reset_peak_memory_stats(device_idx) |
| except RuntimeError: |
| pass |
|
|
| 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, |
| use_cache=not args.no_cache, |
| cache_implementation=args.cache_implementation, |
| ) |
| validate_model_device_placement( |
| model, |
| allow_cpu_offload=args.allow_cpu_offload, |
| requested_gpu_count=requested_gpu_count, |
| ) |
|
|
| if args.mode == "atomic": |
| keep_mask = atomic_mask_from_scores(scores, ratio) |
| apply_atomic_mask_to_model(model, keep_mask) |
| mask_summary = { |
| "mode": "atomic", |
| "ratio": ratio, |
| "mask_shape": list(keep_mask.shape), |
| "kept": int(keep_mask.sum()), |
| "pruned": int((~keep_mask).sum()), |
| } |
| else: |
| keep_mask = group_mask_from_scores(scores, ratio) |
| apply_group_mask_to_model( |
| model, |
| keep_mask, |
| group_width=args.group_width, |
| group_indices=group_indices, |
| ) |
| mask_summary = { |
| "mode": "group", |
| "ratio": ratio, |
| "mask_shape": list(keep_mask.shape), |
| "kept_groups": int(keep_mask.sum()), |
| "pruned_groups": int((~keep_mask).sum()), |
| "group_width": int(args.group_width), |
| "group_indices_shape": list(group_indices.shape) if group_indices is not None else None, |
| } |
|
|
| metrics = evaluate_token_cache( |
| model, |
| args.cache_path, |
| batch_size=args.batch_size, |
| max_chunks=args.max_chunks, |
| use_cache=not args.no_cache, |
| cache_implementation=args.cache_implementation, |
| ) |
| metrics["mask_summary"] = mask_summary |
| output_path = output_dir / f"{args.mode}_pruned_{int(ratio * 100):02d}pct_loss.json" |
| write_json(output_path, metrics) |
| print(output_path) |
| del model |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
