Add Reasoning-Aware Compression (RAC) pruning example for reasoning LLMs

[PKUWZP]

What

Adds compression/reasoning_aware_compression/, an example implementing Reasoning-Aware Compression (RAC) from Reasoning Models Can be Accurately Pruned Via Chain-of-Thought Reconstruction (ICLR 2026). Authors' reference code: RyanLucas3/Reasoning-Aware-Compression.

Why

Layer-wise pruners (SparseGPT, Wanda, ALPS) minimise ||W X - W' X||_F^2 over calibration activations collected from prompts. Reasoning models invert the usual token economics: the prompt is a few hundred tokens and the chain-of-thought answering it is thousands, so pruning optimises the wrong distribution. The failure mode is worse than a plain accuracy drop — the pruned model rambles, so it gets slower as it gets sparser. In the paper, DeepSeek-R1-Distill-Qwen-7B pruned to 50% with C4 calibration takes 135 min on MATH-500 vs 23 min dense, and accuracy falls 0.936 → 0.744.

RAC changes only the calibration set: it collects the dense model's on-policy CoT traces and concatenates their activations with the prompt activations, so each layer is reconstructed against the activations it will actually see while decoding. Same 7B model at 50%: 0.900 accuracy in 35 min. No retraining, no distillation, drop-in for any layer-wise pruner.

What's in the PR

compression/reasoning_aware_compression/
├── collect_traces.py                    # phase I: on-policy CoT rollouts -> jsonl
├── prune.py                             # phase II: one-shot pruning -> sparse checkpoint
├── rac/{calibration,sequential,sparsegpt,wanda,magnitude,modelutils}.py
├── bash_script/                         # end-to-end run, calibration ablation, ZeRO-Inference traces, lighteval
├── tests/                               # CPU-only pytest
└── requirements.txt

• Phase I — collect_traces.py: rollouts from the dense model with vLLM, plain HF generate, or DeepSpeed ZeRO-Inference (--backend deepspeed: ZeRO-3 parameter sharding with CPU/NVMe offload, following inference/huggingface/zero_inference) so 32B/70B dense models can produce traces on a couple of GPUs.
• Phase II — prune.py: SparseGPT / Wanda / magnitude against --calibration {c4,prompt,rac}. The model is streamed through the accelerator one transformer block at a time — peak device memory is one block plus its Hessians, which is what lets 70B prune on a single 80GB GPU as in the paper. Supports unstructured and n:m (2:4) sparsity, --scope mlp, and --layer-thirds 1,3 for the paper's partial-depth throughput setting.
• bash_script/run_calibration_ablation.sh reproduces the paper's core comparison: same model, same pruner, same 1M-token budget, three calibration sets.

Testing

• pytest tests — 23 CPU-only tests, no network or GPU, a few seconds. Cover calibration-window packing, RAC-vs-prompt-only token content, target sparsity for all three pruners, 2:4 patterns, block/scope selection, batched-equals-unbatched calibration, dead input channels, and that the calibration set actually changes the mask.
• Both entry points were exercised end to end on CPU with a tiny randomly-initialised Qwen2: collect_traces.py --backend hf → prune.py --calibration rac → reload the saved checkpoint and verify the sparsity pattern.
• Not run in this PR: the full 1.5B–70B reproductions. The results quoted in the README are the paper's published numbers and are labelled as such.

Notes for reviewers

• rac/sparsegpt.py adapts IST-DASLab/sparsegpt (Apache-2.0), rac/wanda.py adapts locuslab/wanda (MIT); both are attributed in-file and in the README.
• Self-contained: no dependency on the authors' open-r1/TRL fork. Core deps are torch + transformers + datasets + accelerate; vLLM, DeepSpeed and lighteval are optional per phase.
• The pruned checkpoint is a plain Hugging Face directory, so it serves through DeepSpeed-Inference, vLLM or transformers unchanged.

[delock]

It might worth to add a guard to assert generated sequence is consistent among all ranks. Just in case people use a different form of zero3 (i.e. zero3 ranks >=2).

[delock]

might better use 'almost exactly' because the ids reconstributed from teacher decoded text might not be exactly as the teacher hidden states.

[delock]

Some optional suggestions (non-blocking):
1. use_cache not restored on exception (sequential.py:168–169)

use_cache = getattr(model.config, "use_cache", False)
model.config.use_cache = False
# ... pruning ...
model.config.use_cache = use_cache  # not reached if pruning throws

If pruning fails partway, model.config.use_cache stays False. A try/finally would make this robust.

2. n_completion counting comment (collect_traces.py:289)

n_completion = int((sequence != tokenizer.pad_token_id).sum())

Since pad_token = eos_token (L251), the terminal eos is excluded from the completion token count. A one-line comment would prevent future readers from misreading it as "generated length."

3. README "exactly" wording (L57–59)

"Teacher-forcing a sampled trace reproduces exactly the hidden states..."

Traces are stored as decoded text and re-tokenized in calibration.py:150, so a small number of tokens at document boundaries may differ from the original generation. Impact on calibration over 1M tokens is negligible, but
"closely approximates" would be more accurate than "exactly."

4. ZeRO-Inference loop termination (collect_traces.py:295)

if total >= args.trace_tokens:
    break

total is accumulated per-rank from generated content. On a homogeneous cluster this holds in practice, but cross-rank floating-point non-determinism could theoretically cause divergence → deadlock. An optional dist. all_reduce on total before the check would eliminate this latent risk.

5. Empty find_linear_layers warning (modelutils.py:61–78)

For unsupported architectures (e.g., GPT-2/BLOOM using Conv1D), find_linear_layers returns an empty dict and pruning silently produces a "dense" checkpoint. A warning would help catch this edge case.

[tohtana]

Look good to me. I left a comment about the license file of Wanda.

[tohtana]

I think we need to include the license file.