Add Reasoning-Aware Compression (RAC) pruning recipe for reasoning models (#32414)
Co-authored-by: Ryan Lucas <ryanluc@mit.edu> Co-authored-by: Kayhan Behdin <kbehdin@linkedin.com> Co-authored-by: Zhipeng Wang <zwanga@wustl.edu>
This commit is contained in:
Zhipeng Wang
co-authored by
Ryan Lucas
Kayhan Behdin
Zhipeng Wang
Ryan Lucas
Kayhan Behdin
Zhipeng Wang
parent
03c1d58112
commit
bfb224ff01
+232
@@ -0,0 +1,232 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Phase II of Reasoning-Aware Compression (RAC): one-shot prune a reasoning model
|
||||
against the on-policy chain-of-thought calibration set built by
|
||||
`rac_collect_traces.py`.
|
||||
|
||||
RAC (https://arxiv.org/abs/2509.12464, ICLR 2026) does not change the pruning
|
||||
solver. Its whole contribution is which activations the solver reconstructs:
|
||||
prompt tokens *plus* the model's own decode tokens (paper Eq. 7) instead of
|
||||
prompt tokens alone. So this script is deliberately thin -- it hands the RAC
|
||||
calibration set to `llm-compressor`'s SparseGPT or Wanda implementation and
|
||||
saves the result as a checkpoint SGLang can serve.
|
||||
|
||||
Requires `llm-compressor`, which is NOT an SGLang dependency:
|
||||
|
||||
pip install "llmcompressor>=0.12.0"
|
||||
|
||||
Example (reproduces the paper's 50%-sparsity math setting):
|
||||
|
||||
python rac_prune.py \
|
||||
--model-path deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B \
|
||||
--calibration ./rac_traces_math/traces.jsonl \
|
||||
--sparsity 0.5 \
|
||||
--output-dir ./rac_pruned
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
from typing import List, Optional
|
||||
|
||||
import torch
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
INSTALL_HINT = (
|
||||
"This script needs llm-compressor, which SGLang does not depend on.\n"
|
||||
"Install it with:\n\n"
|
||||
' pip install "llmcompressor>=0.12.0"\n'
|
||||
)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="One-shot prune a reasoning model on RAC calibration traces.",
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
|
||||
)
|
||||
|
||||
parser.add_argument("--model-path", required=True, help="Dense model to prune.")
|
||||
parser.add_argument(
|
||||
"--calibration",
|
||||
required=True,
|
||||
help="traces.jsonl produced by rac_collect_traces.py.",
|
||||
)
|
||||
parser.add_argument("--output-dir", required=True)
|
||||
|
||||
parser.add_argument(
|
||||
"--method",
|
||||
choices=["sparsegpt", "wanda"],
|
||||
default="sparsegpt",
|
||||
help="Layer-wise solver. The paper's headline results use SparseGPT.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--sparsity",
|
||||
type=float,
|
||||
default=0.5,
|
||||
help="Layer-wise sparsity. The paper sweeps 0.2-0.5.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--mask-structure",
|
||||
default="0:0",
|
||||
help="'0:0' is unstructured (the paper's setting). '2:4' gives a "
|
||||
"semi-structured mask.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-seq-length",
|
||||
type=int,
|
||||
default=8192,
|
||||
help="Calibration sequences longer than this are truncated.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-samples",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Cap on calibration sequences used. Default: use all of them.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pipeline",
|
||||
default="sequential",
|
||||
help="llm-compressor calibration pipeline. 'sequential' keeps only one "
|
||||
"decoder layer's Hessians resident, which is what fits on one GPU.",
|
||||
)
|
||||
parser.add_argument("--dtype", default="bfloat16")
|
||||
parser.add_argument("--device-map", default="auto")
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def load_calibration_sequences(
|
||||
*, path: str, max_seq_length: int, num_samples: Optional[int], seed: int
|
||||
) -> List[List[int]]:
|
||||
"""Read RAC traces back as token id sequences."""
|
||||
sequences = []
|
||||
with open(path, "r", encoding="utf-8") as trace_file:
|
||||
for line in trace_file:
|
||||
input_ids = json.loads(line)["input_ids"]
|
||||
if input_ids:
|
||||
sequences.append(input_ids[:max_seq_length])
|
||||
|
||||
if not sequences:
|
||||
raise ValueError(f"No calibration sequences found in {path}")
|
||||
|
||||
if num_samples is not None and num_samples < len(sequences):
|
||||
random.Random(seed).shuffle(sequences)
|
||||
sequences = sequences[:num_samples]
|
||||
|
||||
return sequences
|
||||
|
||||
|
||||
def build_calibration_dataloader(sequences: List[List[int]]) -> DataLoader:
|
||||
"""Wrap token id sequences as batches llm-compressor can calibrate on.
|
||||
|
||||
Batch size is 1 on purpose. Batching sequences of different lengths would
|
||||
require padding, and pad-token activations would enter the layer-wise
|
||||
Hessian as if they were real ones -- exactly the calibration contamination
|
||||
RAC is about avoiding.
|
||||
"""
|
||||
|
||||
def collate(batch: List[List[int]]) -> dict:
|
||||
input_ids = torch.tensor(batch[0], dtype=torch.long).unsqueeze(0)
|
||||
return {
|
||||
"input_ids": input_ids,
|
||||
"attention_mask": torch.ones_like(input_ids),
|
||||
}
|
||||
|
||||
return DataLoader(sequences, batch_size=1, shuffle=False, collate_fn=collate)
|
||||
|
||||
|
||||
def build_recipe(*, method: str, sparsity: float, mask_structure: str):
|
||||
"""Instantiate the layer-wise solver. RAC leaves this untouched."""
|
||||
try:
|
||||
from llmcompressor.modifiers.pruning import (
|
||||
SparseGPTModifier,
|
||||
WandaPruningModifier,
|
||||
)
|
||||
except ImportError as exc:
|
||||
raise ImportError(INSTALL_HINT) from exc
|
||||
|
||||
modifier_cls = SparseGPTModifier if method == "sparsegpt" else WandaPruningModifier
|
||||
return modifier_cls(
|
||||
sparsity=sparsity,
|
||||
mask_structure=mask_structure,
|
||||
targets=["Linear"],
|
||||
ignore=["re:.*lm_head"],
|
||||
)
|
||||
|
||||
|
||||
def measure_sparsity(model) -> float:
|
||||
"""Fraction of zeros across the pruned Linear weights."""
|
||||
num_zeros = 0
|
||||
num_weights = 0
|
||||
for name, module in model.named_modules():
|
||||
if not isinstance(module, torch.nn.Linear) or "lm_head" in name:
|
||||
continue
|
||||
weight = module.weight
|
||||
num_zeros += int((weight == 0).sum().item())
|
||||
num_weights += weight.numel()
|
||||
|
||||
return num_zeros / num_weights if num_weights else 0.0
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
|
||||
try:
|
||||
from llmcompressor import oneshot
|
||||
except ImportError as exc:
|
||||
raise ImportError(INSTALL_HINT) from exc
|
||||
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
sequences = load_calibration_sequences(
|
||||
path=args.calibration,
|
||||
max_seq_length=args.max_seq_length,
|
||||
num_samples=args.num_samples,
|
||||
seed=args.seed,
|
||||
)
|
||||
num_calibration_tokens = sum(len(sequence) for sequence in sequences)
|
||||
print(
|
||||
f"[rac] calibrating on {len(sequences)} sequences "
|
||||
f"({num_calibration_tokens} tokens) from {args.calibration}"
|
||||
)
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
args.model_path,
|
||||
dtype=getattr(torch, args.dtype),
|
||||
device_map=args.device_map,
|
||||
)
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.model_path)
|
||||
|
||||
print(
|
||||
f"[rac] pruning to {args.sparsity:.0%} sparsity "
|
||||
f"with {args.method} (mask_structure={args.mask_structure})"
|
||||
)
|
||||
model = oneshot(
|
||||
model=model,
|
||||
processor=tokenizer,
|
||||
dataset=build_calibration_dataloader(sequences),
|
||||
recipe=build_recipe(
|
||||
method=args.method,
|
||||
sparsity=args.sparsity,
|
||||
mask_structure=args.mask_structure,
|
||||
),
|
||||
pipeline=args.pipeline,
|
||||
output_dir=args.output_dir,
|
||||
)
|
||||
|
||||
realized = measure_sparsity(model)
|
||||
print(f"\n[rac] realized sparsity: {realized:.2%} (target {args.sparsity:.2%})")
|
||||
print(f"[rac] checkpoint written to {os.path.abspath(args.output_dir)}")
|
||||
print("\nServe it:")
|
||||
print(f" python -m sglang.launch_server --model-path {args.output_dir}")
|
||||
print("\nOr score it against the dense model:")
|
||||
print(
|
||||
f" python rac_serve_and_eval.py --model-path {args.output_dir} "
|
||||
"--num-problems 100"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user