| |
| """Run MMLU against a loaded baseline, atomic-pruned, or repacked grouped model.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import sys |
| from contextlib import nullcontext |
| from pathlib import Path |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) |
|
|
| import numpy as np |
|
|
| from heapr.grouped_model import RepackedExpandedGroupedLagunaContext |
| from heapr.model_utils import ( |
| build_max_memory, |
| load_causal_lm, |
| load_tokenizer, |
| validate_model_device_placement, |
| ) |
| from heapr.prune import apply_atomic_mask_to_model, apply_group_mask_to_model, atomic_mask_from_scores |
| from heapr.utils import require_torch |
|
|
|
|
| def parse_args(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--model-id", required=True) |
| parser.add_argument("--output-path", required=True) |
| parser.add_argument( |
| "--mode", |
| choices=["baseline", "atomic", "native-group", "repacked-group"], |
| default="baseline", |
| ) |
| parser.add_argument("--scores-path") |
| parser.add_argument("--ratio", type=float) |
| parser.add_argument("--group-mask") |
| parser.add_argument("--group-indices") |
| parser.add_argument( |
| "--child-budget-mode", |
| choices=["backfill", "layer-scaled", "parent-weighted"], |
| default="backfill", |
| ) |
| parser.add_argument("--tasks", default="mmlu") |
| parser.add_argument("--limit", type=float) |
| parser.add_argument("--num-fewshot", type=int) |
| parser.add_argument("--max-gen-toks", type=int) |
| parser.add_argument("--batch-size", default="1") |
| parser.add_argument("--revision") |
| 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") |
| return parser.parse_args() |
|
|
|
|
| def _batch_size(value: str): |
| return int(value) if value.isdigit() else value |
|
|
|
|
| def _json_default(value): |
| if isinstance(value, np.generic): |
| return value.item() |
| if isinstance(value, np.ndarray): |
| return value.tolist() |
| return str(value) |
|
|
|
|
| def _generation_kwargs(max_gen_toks: int | None): |
| if max_gen_toks is None: |
| return None |
| return {"max_gen_toks": max_gen_toks} |
|
|
|
|
| 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 == "atomic" and (args.scores_path is None or args.ratio is None): |
| raise ValueError("--scores-path and --ratio are required for atomic mode") |
| if args.mode in {"native-group", "repacked-group"} and ( |
| args.group_mask is None or args.group_indices is None |
| ): |
| raise ValueError("--group-mask and --group-indices are required for grouped modes") |
|
|
| 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, |
| use_cache=False, |
| ) |
| if args.max_gen_toks is not None: |
| |
| |
| model.generation_config.max_new_tokens = args.max_gen_toks |
| validate_model_device_placement( |
| model, |
| allow_cpu_offload=args.allow_cpu_offload, |
| requested_gpu_count=requested_gpu_count, |
| ) |
| tokenizer = load_tokenizer(args.model_id, revision=args.revision) |
|
|
| context = nullcontext(model) |
| pruning_summary = {"mode": args.mode} |
| if args.mode == "atomic": |
| keep_mask = atomic_mask_from_scores(np.load(args.scores_path), args.ratio) |
| apply_atomic_mask_to_model(model, keep_mask) |
| pruning_summary.update( |
| ratio=args.ratio, |
| kept=int(keep_mask.sum()), |
| pruned=int((~keep_mask).sum()), |
| ) |
| elif args.mode == "native-group": |
| keep_mask = np.load(args.group_mask) |
| group_indices = np.load(args.group_indices) |
| apply_group_mask_to_model(model, keep_mask, group_indices=group_indices) |
| pruning_summary.update( |
| kept_groups=int(keep_mask.sum()), |
| pruned_groups=int((~keep_mask).sum()), |
| ) |
| elif args.mode == "repacked-group": |
| keep_mask = np.load(args.group_mask) |
| group_indices = np.load(args.group_indices) |
| context = RepackedExpandedGroupedLagunaContext( |
| model, |
| keep_group_masks=keep_mask, |
| group_indices=group_indices, |
| child_budget_mode=args.child_budget_mode, |
| ) |
| pruning_summary.update( |
| kept_groups=int(keep_mask.sum()), |
| pruned_groups=int((~keep_mask).sum()), |
| child_budget_mode=args.child_budget_mode, |
| ) |
|
|
| from lm_eval import evaluator |
| from lm_eval.models.huggingface import HFLM |
|
|
| with context: |
| harness_model = HFLM( |
| pretrained=model, |
| tokenizer=tokenizer, |
| backend="causal", |
| batch_size=_batch_size(args.batch_size), |
| ) |
| results = evaluator.simple_evaluate( |
| model=harness_model, |
| tasks=[task.strip() for task in args.tasks.split(",") if task.strip()], |
| num_fewshot=args.num_fewshot, |
| batch_size=_batch_size(args.batch_size), |
| limit=args.limit, |
| gen_kwargs=_generation_kwargs(args.max_gen_toks), |
| ) |
|
|
| payload = { |
| "model_id": args.model_id, |
| "pruning": pruning_summary, |
| "tasks": args.tasks, |
| "limit": args.limit, |
| "num_fewshot": args.num_fewshot, |
| "max_gen_toks": args.max_gen_toks, |
| "batch_size": args.batch_size, |
| "results": results, |
| } |
| output_path = Path(args.output_path) |
| output_path.parent.mkdir(parents=True, exist_ok=True) |
| output_path.write_text( |
| json.dumps(payload, indent=2, sort_keys=True, default=_json_default) + "\n", |
| encoding="utf-8", |
| ) |
| print(output_path) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
