From a801dbef0e07c217565403524ee72508af451192 Mon Sep 17 00:00:00 2001 From: Zhipeng Wang Date: Sat, 25 Jul 2026 22:53:39 -0700 Subject: [PATCH 1/2] Add Reasoning-Aware Compression (RAC) pruning example Implements "Reasoning Models Can be Accurately Pruned Via Chain-of-Thought Reconstruction" (Lucas et al., ICLR 2026, arXiv:2509.12464) under compression/reasoning_aware_compression. Layer-wise pruners calibrate on prompt activations, which is the wrong distribution for reasoning models: decoding dominates their token budget, so pruned models both lose accuracy and get slower by rambling. RAC calibrates on the dense model's own on-policy chain-of-thought instead, as a drop-in change to SparseGPT/Wanda. Two phases: collect_traces.py on-policy CoT rollouts -> jsonl, with vLLM, plain HF, or DeepSpeed ZeRO-Inference (ZeRO-3 with CPU/NVMe parameter offload) for 32B/70B dense models prune.py one-shot SparseGPT / Wanda / magnitude pruning against a c4, prompt-only, or RAC calibration set Pruning streams the model through the accelerator one transformer block at a time, so peak device memory is one block plus its Hessians. Supports unstructured and n:m sparsity, MLP-only scope, and per-third block selection. Includes bash scripts for the end-to-end run, the c4/prompt/RAC ablation and lighteval MATH-500 evaluation, plus CPU-only unit tests (pytest tests, 23 tests, no network or GPU). Signed-off-by: Zhipeng Wang --- compression/README.md | 7 + .../reasoning_aware_compression/README.md | 220 +++++++++++ .../collect_traces_zero_inference.sh | 44 +++ .../bash_script/eval_math500.sh | 39 ++ .../bash_script/run_calibration_ablation.sh | 55 +++ .../bash_script/run_rac_qwen1_5b.sh | 56 +++ .../collect_traces.py | 355 ++++++++++++++++++ .../reasoning_aware_compression/prune.py | 218 +++++++++++ .../rac/__init__.py | 38 ++ .../rac/calibration.py | 257 +++++++++++++ .../rac/magnitude.py | 52 +++ .../rac/modelutils.py | 145 +++++++ .../rac/sequential.py | 266 +++++++++++++ .../rac/sparsegpt.py | 168 +++++++++ .../reasoning_aware_compression/rac/wanda.py | 73 ++++ .../requirements.txt | 15 + .../tests/test_calibration.py | 131 +++++++ .../tests/test_pruning.py | 192 ++++++++++ 18 files changed, 2331 insertions(+) create mode 100644 compression/reasoning_aware_compression/README.md create mode 100755 compression/reasoning_aware_compression/bash_script/collect_traces_zero_inference.sh create mode 100755 compression/reasoning_aware_compression/bash_script/eval_math500.sh create mode 100755 compression/reasoning_aware_compression/bash_script/run_calibration_ablation.sh create mode 100755 compression/reasoning_aware_compression/bash_script/run_rac_qwen1_5b.sh create mode 100644 compression/reasoning_aware_compression/collect_traces.py create mode 100644 compression/reasoning_aware_compression/prune.py create mode 100644 compression/reasoning_aware_compression/rac/__init__.py create mode 100644 compression/reasoning_aware_compression/rac/calibration.py create mode 100644 compression/reasoning_aware_compression/rac/magnitude.py create mode 100644 compression/reasoning_aware_compression/rac/modelutils.py create mode 100644 compression/reasoning_aware_compression/rac/sequential.py create mode 100644 compression/reasoning_aware_compression/rac/sparsegpt.py create mode 100644 compression/reasoning_aware_compression/rac/wanda.py create mode 100644 compression/reasoning_aware_compression/requirements.txt create mode 100644 compression/reasoning_aware_compression/tests/test_calibration.py create mode 100644 compression/reasoning_aware_compression/tests/test_pruning.py diff --git a/compression/README.md b/compression/README.md index bcaa937c4..63b07c17d 100644 --- a/compression/README.md +++ b/compression/README.md @@ -3,3 +3,10 @@ Examples in this folder are helpful to try out some features and models that take advantage of the DeepSpeed compression library. A detailed tutorial for understanding and using DeepSpeed model compression features can be seen from here: https://www.deepspeed.ai/tutorials/model-compression/ + +| Example | Description | +| --- | --- | +| [bert](bert) | Quantization, pruning and layer reduction on BERT (ZeroQuant, XTC) | +| [gpt2](gpt2) | ZeroQuant post-training quantization on GPT-2 | +| [cifar](cifar) | Channel pruning and quantization on a CIFAR ResNet | +| [reasoning_aware_compression](reasoning_aware_compression) | One-shot pruning of reasoning LLMs (DeepSeek-R1 distills, Qwen3) calibrated on their own chain-of-thought traces — [RAC, ICLR 2026](https://arxiv.org/abs/2509.12464) | diff --git a/compression/reasoning_aware_compression/README.md b/compression/reasoning_aware_compression/README.md new file mode 100644 index 000000000..de5cb8b11 --- /dev/null +++ b/compression/reasoning_aware_compression/README.md @@ -0,0 +1,220 @@ +# Reasoning-Aware Compression (RAC) + +One-shot pruning of reasoning LLMs, calibrated on the model's **own chain-of-thought**. + +Implements [*Reasoning Models Can be Accurately Pruned Via Chain-of-Thought +Reconstruction*](https://arxiv.org/abs/2509.12464) (Lucas, Behdin, Wang, Tang, Song, +Mazumder — ICLR 2026). Reference code from the authors: +[RyanLucas3/Reasoning-Aware-Compression](https://github.com/RyanLucas3/Reasoning-Aware-Compression). + +## Why + +Layer-wise pruning methods (SparseGPT, Wanda, ALPS) minimise a reconstruction error + +``` +min ||W X - W' X||_F^2 s.t. ||W'||_0 <= S +``` + +over calibration activations `X`, which are traditionally collected by running **prompts** +(C4 documents, task inputs) through the dense model. That is a reasonable proxy when +inference is prompt-dominated. Reasoning models are the opposite: a MATH-500 problem is a +few hundred tokens and the chain-of-thought that answers it is thousands. Pruning +therefore optimises the wrong distribution, and the failure 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 +versus 23 min dense, while accuracy falls from 0.936 to 0.744. + +RAC changes the calibration set and nothing else: it collects the dense model's on-policy +CoT traces and concatenates their activations with the prompt activations, + +``` +X_RAC = [ X_prompt | X_decode ] +``` + +so each layer is reconstructed against the activations it will really see while decoding. +It is a drop-in change to any layer-wise pruner, needs no retraining or distillation, and +recovers most of the loss: the same 7B model at 50% sparsity reaches 0.900 accuracy in +35 min. + +## How it works + +``` + phase I: collect_traces.py phase II: prune.py + ┌────────────────────────────────────┐ ┌──────────────────────────────────────┐ + │ task prompts (math / code) │ │ calibration windows │ + │ │ │ │ c4 | prompt | rac │ + │ ▼ │ │ │ │ + │ dense reasoning model │ │ ▼ │ + │ vLLM | HF | ZeRO-Inference │ │ block 0 → prune → propagate │ + │ │ │ │ block 1 → prune → propagate │ + │ ▼ │ │ ... (one block on GPU at a time)│ + │ prompt + on-policy CoT ──── jsonl ├─────▶│ block L → prune │ + └────────────────────────────────────┘ │ │ │ + │ ▼ sparse checkpoint │ + └──────────────────────────────────────┘ +``` + +Teacher-forcing a sampled trace reproduces exactly the hidden states the model computed +while generating it, so phase II recovers the decode-time activations of Algorithm 1 in +the paper without re-running generation during pruning. + +## Layout + +``` +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.py # c4 / prompt / rac calibration sets <- the method +│ ├── sequential.py # block-by-block calibration driver +│ ├── sparsegpt.py # SparseGPT Hessian + reconstruction +│ ├── wanda.py # Wanda activation-norm pruner +│ ├── magnitude.py # magnitude baseline (calibration-free) +│ └── modelutils.py # layer discovery, sparsity reporting +├── bash_script/ +│ ├── run_rac_qwen1_5b.sh # end-to-end: traces -> prune -> eval +│ ├── run_calibration_ablation.sh # c4 vs prompt vs rac, same budget +│ ├── collect_traces_zero_inference.sh # traces for 32B/70B via ZeRO-Inference +│ └── eval_math500.sh # lighteval MATH-500 (accuracy + runtime) +├── tests/ # CPU-only unit tests (pytest) +└── requirements.txt +``` + +## Quick start + +```bash +cd compression/reasoning_aware_compression +pip install -r requirements.txt +pip install vllm # optional but much faster for phase I +``` + +**Phase I — collect on-policy traces** (paper: 1M tokens, `T_max = 8192`): + +```bash +python collect_traces.py \ + --model deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B \ + --dataset open-r1/OpenR1-Math-220k --prompt-column problem \ + --trace-tokens 1_000_000 --max-new-tokens 8192 \ + --output outputs/traces/r1-qwen-1.5b_math.jsonl +``` + +**Phase II — prune against them**: + +```bash +python prune.py \ + --model deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B \ + --calibration rac --traces outputs/traces/r1-qwen-1.5b_math.jsonl \ + --sparsity 0.5 --nsamples 512 --seqlen 2048 \ + --output outputs/checkpoints/r1-qwen-1.5b-rac-50 +``` + +**Evaluate** (accuracy *and* runtime — both matter here): + +```bash +MODEL_DIR=outputs/checkpoints/r1-qwen-1.5b-rac-50 bash bash_script/eval_math500.sh +``` + +Or run all three phases at once with `bash bash_script/run_rac_qwen1_5b.sh`. + +### Calibration sets + +`--calibration` selects the ablation axis of the paper. Everything else stays fixed. + +| Flag | Calibration tokens | Notes | +| --- | --- | --- | +| `--calibration c4` | generic web text | SparseGPT/Wanda default; `--dataset allenai/c4 --dataset-config en` | +| `--calibration prompt` | task prompts only | pass `--traces` to reuse exactly the prompts RAC saw, or `--dataset` for a fresh set | +| `--calibration rac` | prompts **+** on-policy CoT | the method; requires `--traces` | + +`--nsamples × --seqlen` is the token budget (`512 × 2048 = 1M`, as in the paper). +Documents are packed into fixed-length windows, so no padding tokens enter the Hessian. + +### Other knobs + +| Flag | Meaning | +| --- | --- | +| `--pruning-method` | `sparsegpt` (default), `wanda`, `magnitude`. RAC helps all calibrated methods (Table 9). | +| `--prune-n 2 --prune-m 4` | semi-structured 2:4 sparsity for Ampere+ sparse tensor cores | +| `--scope mlp` | prune only feed-forward projections, leave attention dense | +| `--layer-thirds 1,3` | prune only the first and last third of the decoder stack (Table 6) | +| `--batch-size` | calibration samples per forward pass; raise it if the GPU has headroom | + +## Scaling to large models + +**Pruning** streams the model through the accelerator one transformer block at a time — +the block, its Hessians and one activation slice are the only device-resident tensors — +so a 70B checkpoint prunes on a single 80GB GPU, as in the paper. Nothing needs to be +sharded and `prune.py` is a single-process script. Host memory holds the weights plus two +activation buffers of `nsamples × seqlen × hidden_size` (~17GB each at the default budget +for a 70B model in bf16); drop `--nsamples` if that is tight. + +**Trace collection** is the expensive phase and is where DeepSpeed helps. For models that +do not fit in aggregate GPU memory, `--backend deepspeed` runs generation under +ZeRO-Inference (ZeRO-3 parameter sharding with CPU or NVMe offload), following +[`inference/huggingface/zero_inference`](../../inference/huggingface/zero_inference): + +```bash +NUM_GPUS=2 MODEL=deepseek-ai/DeepSeek-R1-Distill-Qwen-32B \ + bash bash_script/collect_traces_zero_inference.sh +``` + +Use `NVME_OFFLOAD_DIR=/path/on/nvme` to spill parameters to NVMe instead of CPU memory. +When the model does fit, `--backend vllm` is considerably faster. + +The pruned checkpoint is a standard Hugging Face directory, so it serves through +DeepSpeed-Inference, vLLM or plain `transformers` unchanged. Unstructured sparsity gives +memory and (with a sparse kernel) compute savings; 2:4 checkpoints run on Ampere+ sparse +tensor cores — the paper measures 1561 tok/s at 2:4 versus 1426 tok/s dense for +DeepSeek-R1-Distill-14B. + +## Results from the paper + +MATH-500 acc@1 and total evaluation time, one-shot SparseGPT, 1M calibration tokens, +32k decode budget. These are the authors' published numbers (Tables 1 and 2), reproduced +here for reference — this example is the pipeline that generates them, not a re-run. + +| Model | Sparsity | C4 | Prompt-only | **RAC** | Runtime C4 → RAC | +| --- | --- | --- | --- | --- | --- | +| R1-Distill-Qwen-1.5B (dense 0.832) | 40% | 0.658 | 0.728 | **0.774** | 58 → 32 min | +| | 50% | 0.356 | 0.496 | **0.664** | 157 → 57 min | +| R1-Distill-Qwen-7B (dense 0.936) | 40% | 0.890 | 0.898 | **0.912** | 38 → 29 min | +| | 50% | 0.744 | 0.812 | **0.900** | 135 → 35 min | +| Qwen3-8B (dense 0.962) | 40% | 0.906 | 0.944 | **0.968** | 58 → 29 min | +| | 50% | 0.564 | 0.470 | **0.862** | 259 → 17 min | + +The gap widens with sparsity, and the runtime column is the practical point: the C4- +calibrated models are both less accurate and slower than the dense baseline. The paper +reports the same pattern on AIME-25 and LiveCodeBench, and finds that on-policy traces +(from the model being pruned) beat off-policy traces from a larger model. + +## Tests + +CPU-only, no network, a few seconds: + +```bash +cd compression/reasoning_aware_compression +pytest tests +``` + +They cover calibration-window packing, the RAC-vs-prompt distinction, target sparsity for +SparseGPT/Wanda/magnitude, 2:4 patterns, block and scope selection, and batched-equals- +unbatched calibration on a tiny randomly-initialised Qwen2. + +## Attribution + +- `rac/sparsegpt.py` adapts [IST-DASLab/sparsegpt](https://github.com/IST-DASLab/sparsegpt) (Apache-2.0). +- `rac/wanda.py` adapts [locuslab/wanda](https://github.com/locuslab/wanda) (MIT). +- The RAC method and its reference implementation are by the paper's authors + ([repo](https://github.com/RyanLucas3/Reasoning-Aware-Compression)); this example is a + standalone reimplementation for DeepSpeedExamples. + +```bibtex +@inproceedings{lucas2026rac, + title = {Reasoning Models Can be Accurately Pruned Via Chain-of-Thought Reconstruction}, + author = {Lucas, Ryan and Behdin, Kayhan and Wang, Zhipeng and Tang, Shao and + Song, Qingquan and Mazumder, Rahul}, + booktitle = {International Conference on Learning Representations (ICLR)}, + year = {2026}, + url = {https://arxiv.org/abs/2509.12464} +} +``` diff --git a/compression/reasoning_aware_compression/bash_script/collect_traces_zero_inference.sh b/compression/reasoning_aware_compression/bash_script/collect_traces_zero_inference.sh new file mode 100755 index 000000000..3b2e18493 --- /dev/null +++ b/compression/reasoning_aware_compression/bash_script/collect_traces_zero_inference.sh @@ -0,0 +1,44 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +# +# Trace collection for models that do not fit in aggregate GPU memory, using +# ZeRO-Inference: ZeRO-3 shards the parameters across ranks and streams them +# from CPU (or NVMe) as each layer runs. All ranks generate the same batch in +# lockstep (synced_gpus) and rank 0 writes the JSONL. +# +# See inference/huggingface/zero_inference/README.md in this repo for the +# memory/throughput characteristics of the offload paths. +# +# cd compression/reasoning_aware_compression +# NUM_GPUS=2 MODEL=deepseek-ai/DeepSeek-R1-Distill-Qwen-32B \ +# bash bash_script/collect_traces_zero_inference.sh +set -euo pipefail + +NUM_GPUS=${NUM_GPUS:-2} +MODEL=${MODEL:-deepseek-ai/DeepSeek-R1-Distill-Qwen-32B} +DATASET=${DATASET:-open-r1/OpenR1-Math-220k} +PROMPT_COLUMN=${PROMPT_COLUMN:-problem} +TRACE_TOKENS=${TRACE_TOKENS:-1000000} +BATCH_SIZE=${BATCH_SIZE:-4} +OUTPUT=${OUTPUT:-./outputs/traces/$(basename "${MODEL}")_math.jsonl} +# Set NVME_OFFLOAD_DIR to spill parameters to NVMe instead of CPU memory. +NVME_OFFLOAD_DIR=${NVME_OFFLOAD_DIR:-} + +OFFLOAD_ARGS=(--cpu-offload) +if [[ -n "${NVME_OFFLOAD_DIR}" ]]; then + OFFLOAD_ARGS=(--nvme-offload-dir "${NVME_OFFLOAD_DIR}") +fi + +deepspeed --num_gpus "${NUM_GPUS}" collect_traces.py \ + --backend deepspeed \ + --model "${MODEL}" \ + --dataset "${DATASET}" \ + --prompt-column "${PROMPT_COLUMN}" \ + --trace-tokens "${TRACE_TOKENS}" \ + --max-new-tokens 8192 \ + --batch-size "${BATCH_SIZE}" \ + --output "${OUTPUT}" \ + "${OFFLOAD_ARGS[@]}" + +echo "[traces] written to ${OUTPUT}" diff --git a/compression/reasoning_aware_compression/bash_script/eval_math500.sh b/compression/reasoning_aware_compression/bash_script/eval_math500.sh new file mode 100755 index 000000000..eb5046ef5 --- /dev/null +++ b/compression/reasoning_aware_compression/bash_script/eval_math500.sh @@ -0,0 +1,39 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +# +# MATH-500 acc@1 for a pruned (or dense) reasoning checkpoint, following the +# open-r1 evaluation pipeline used in the RAC paper. Reports the total runtime +# too: heavily pruned reasoning models ramble, and the paper's headline is that +# RAC keeps both accuracy *and* decode length close to dense. +# +# pip install "lighteval[math,vllm]" +# cd compression/reasoning_aware_compression +# MODEL_DIR=outputs/checkpoints/... bash bash_script/eval_math500.sh +set -euo pipefail + +MODEL_DIR=${MODEL_DIR:?set MODEL_DIR to the checkpoint to evaluate} +NUM_GPUS=${NUM_GPUS:-1} +MAX_NEW_TOKENS=${MAX_NEW_TOKENS:-32768} +TASK=${TASK:-"lighteval|math_500|0|0"} +OUTPUT_DIR=${OUTPUT_DIR:-./outputs/eval/$(basename "${MODEL_DIR}")} + +export VLLM_WORKER_MULTIPROC_METHOD=spawn + +MODEL_ARGS="model_name=${MODEL_DIR},\ +dtype=bfloat16,\ +trust_remote_code=true,\ +max_model_length=${MAX_NEW_TOKENS},\ +gpu_memory_utilization=0.8,\ +data_parallel_size=${NUM_GPUS},\ +generation_parameters={max_new_tokens:${MAX_NEW_TOKENS},temperature:0.6,top_p:0.95}" + +START=$SECONDS +lighteval vllm "${MODEL_ARGS}" "${TASK}" \ + --use-chat-template \ + --output-dir "${OUTPUT_DIR}" \ + --save-details +echo "[eval] ${TASK} on ${MODEL_DIR} took $(( (SECONDS - START) / 60 )) min" + +# Coding: "extended|lcb:codegeneration|0|0" (needs lighteval[extended-tasks]). +# AIME-25: "lighteval|aime25|0|0". diff --git a/compression/reasoning_aware_compression/bash_script/run_calibration_ablation.sh b/compression/reasoning_aware_compression/bash_script/run_calibration_ablation.sh new file mode 100755 index 000000000..cc726766a --- /dev/null +++ b/compression/reasoning_aware_compression/bash_script/run_calibration_ablation.sh @@ -0,0 +1,55 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +# +# Reproduces the core comparison of the RAC paper: the same model, the same +# pruning algorithm and the same 1M-token budget, pruned three times with +# three different calibration sets. +# +# c4 generic web text (the SparseGPT default) +# prompt the task prompts only +# rac those prompts plus the dense model's own chain-of-thought +# +# Requires traces from collect_traces.py; run bash_script/run_rac_qwen1_5b.sh +# first (or set TRACES to an existing file). +# +# cd compression/reasoning_aware_compression +# TRACES=outputs/traces/DeepSeek-R1-Distill-Qwen-1.5B_math.jsonl \ +# bash bash_script/run_calibration_ablation.sh +set -euo pipefail + +MODEL=${MODEL:-deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B} +TRACES=${TRACES:?set TRACES to the JSONL written by collect_traces.py} +SPARSITY=${SPARSITY:-0.5} +NSAMPLES=${NSAMPLES:-512} +SEQLEN=${SEQLEN:-2048} +OUTPUT_ROOT=${OUTPUT_ROOT:-./outputs} + +TAG=$(basename "${MODEL}") + +for CALIBRATION in c4 prompt rac; do + OUT=${OUTPUT_ROOT}/checkpoints/${TAG}_${CALIBRATION}_sparsity${SPARSITY} + echo "=== calibration: ${CALIBRATION} -> ${OUT}" + + EXTRA=() + if [[ "${CALIBRATION}" == "c4" ]]; then + # Streams allenai/c4; no traces involved. + EXTRA=(--dataset allenai/c4 --dataset-config en) + else + # Both prompt-only and RAC read the same trace file, so the prompts are + # identical and the only difference is the CoT tokens. + EXTRA=(--traces "${TRACES}") + fi + + python prune.py \ + --model "${MODEL}" \ + --calibration "${CALIBRATION}" \ + --pruning-method sparsegpt \ + --sparsity "${SPARSITY}" \ + --nsamples "${NSAMPLES}" \ + --seqlen "${SEQLEN}" \ + --output "${OUT}" \ + "${EXTRA[@]}" + + echo "MODEL_DIR=${OUT} bash bash_script/eval_math500.sh" +done diff --git a/compression/reasoning_aware_compression/bash_script/run_rac_qwen1_5b.sh b/compression/reasoning_aware_compression/bash_script/run_rac_qwen1_5b.sh new file mode 100755 index 000000000..4d00bd074 --- /dev/null +++ b/compression/reasoning_aware_compression/bash_script/run_rac_qwen1_5b.sh @@ -0,0 +1,56 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +# +# End-to-end RAC on DeepSeek-R1-Distill-Qwen-1.5B: collect on-policy CoT +# traces, prune to 50% with SparseGPT calibrated on them, then evaluate. +# Runs on a single 80GB GPU; trace collection dominates the wall clock. +# +# cd compression/reasoning_aware_compression +# bash bash_script/run_rac_qwen1_5b.sh +set -euo pipefail + +MODEL=${MODEL:-deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B} +DATASET=${DATASET:-open-r1/OpenR1-Math-220k} +PROMPT_COLUMN=${PROMPT_COLUMN:-problem} +SPARSITY=${SPARSITY:-0.5} +TRACE_TOKENS=${TRACE_TOKENS:-1000000} +NSAMPLES=${NSAMPLES:-512} +SEQLEN=${SEQLEN:-2048} +OUTPUT_ROOT=${OUTPUT_ROOT:-./outputs} + +TAG=$(basename "${MODEL}") +TRACES=${OUTPUT_ROOT}/traces/${TAG}_math.jsonl +PRUNED=${OUTPUT_ROOT}/checkpoints/${TAG}_rac_sparsity${SPARSITY} + +# ---------------------------------------------------------------- phase I -- +# On-policy rollouts from the *dense* model. Skipped if the traces exist, so +# the pruning sweep below can be re-run cheaply. +if [[ -s "${TRACES}" ]]; then + echo "[run_rac] reusing traces at ${TRACES}" +else + python collect_traces.py \ + --model "${MODEL}" \ + --dataset "${DATASET}" \ + --prompt-column "${PROMPT_COLUMN}" \ + --trace-tokens "${TRACE_TOKENS}" \ + --max-new-tokens 8192 \ + --temperature 0.6 \ + --top-p 0.95 \ + --output "${TRACES}" +fi + +# --------------------------------------------------------------- phase II -- +python prune.py \ + --model "${MODEL}" \ + --calibration rac \ + --traces "${TRACES}" \ + --pruning-method sparsegpt \ + --sparsity "${SPARSITY}" \ + --nsamples "${NSAMPLES}" \ + --seqlen "${SEQLEN}" \ + --output "${PRUNED}" \ + --save-report + +echo "[run_rac] pruned checkpoint: ${PRUNED}" +echo "[run_rac] evaluate with: MODEL_DIR=${PRUNED} bash bash_script/eval_math500.sh" diff --git a/compression/reasoning_aware_compression/collect_traces.py b/compression/reasoning_aware_compression/collect_traces.py new file mode 100644 index 000000000..720a23ba1 --- /dev/null +++ b/compression/reasoning_aware_compression/collect_traces.py @@ -0,0 +1,355 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""Phase I of RAC: collect on-policy chain-of-thought traces from a dense model. + +The traces are the model's *own* rollouts on task prompts, which is what makes +the calibration activations match what the model computes while decoding. +Output is a JSONL file consumed by ``prune.py --calibration rac``. + +Three generation backends: + +``vllm`` + Fastest single-node option; use it whenever vLLM is installed. +``hf`` + Plain ``transformers.generate``; works anywhere, slow for 8k-token traces. +``deepspeed`` + ZeRO-Inference: ZeRO-3 shards the weights across ranks and can offload them + to CPU or NVMe, so a 32B/70B dense model generates its traces on a small + number of GPUs. Launch with the DeepSpeed launcher, e.g. + ``deepspeed --num_gpus 2 collect_traces.py --backend deepspeed ...``. + +Example:: + + python collect_traces.py \\ + --model deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B \\ + --dataset open-r1/OpenR1-Math-220k --prompt-column problem \\ + --output traces/math_1.5b.jsonl --trace-tokens 1_000_000 +""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +import os +import time +from pathlib import Path +from typing import List + +DEFAULT_SYSTEM_PROMPT = ( + "You are a helpful AI Assistant that provides well-reasoned and detailed responses. " + "You first think about the reasoning process as an internal monologue and then provide " + "the user with the answer.") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + + model = parser.add_argument_group("model") + model.add_argument("--model", required=True, help="Dense reasoning model to roll out.") + model.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "float16", "float32"]) + model.add_argument("--trust-remote-code", action="store_true") + + data = parser.add_argument_group("prompts") + data.add_argument("--dataset", + required=True, + help="HF dataset name or path to a local .jsonl file of prompts.") + data.add_argument("--dataset-config", default=None) + data.add_argument("--dataset-split", default="train") + data.add_argument("--prompt-column", default="problem", + help="'problem' for OpenR1-Math-220k, 'prompt' for Codeforces.") + data.add_argument("--system-prompt", default=DEFAULT_SYSTEM_PROMPT) + data.add_argument("--no-chat-template", action="store_true", + help="Feed raw prompts instead of applying the chat template.") + data.add_argument("--max-prompts", type=int, default=None, + help="Cap on prompts to roll out (the token budget usually binds first).") + + gen = parser.add_argument_group("generation") + gen.add_argument("--backend", default="auto", choices=["auto", "vllm", "hf", "deepspeed"]) + gen.add_argument("--trace-tokens", type=int, default=1_000_000, + help="Stop once this many prompt+completion tokens have been collected. " + "The paper uses 1M.") + gen.add_argument("--max-new-tokens", type=int, default=8192, + help="T_max in Algorithm 1 of the paper.") + gen.add_argument("--num-generations", type=int, default=1, + help="On-policy rollouts per prompt.") + gen.add_argument("--temperature", type=float, default=0.6) + gen.add_argument("--top-p", type=float, default=0.95) + gen.add_argument("--batch-size", type=int, default=8, help="Prompts per generation call.") + gen.add_argument("--seed", type=int, default=0) + + ds = parser.add_argument_group("deepspeed backend") + ds.add_argument("--tensor-parallel-size", type=int, default=1, help="vLLM tensor parallelism.") + ds.add_argument("--gpu-memory-utilization", type=float, default=0.85, help="vLLM only.") + ds.add_argument("--cpu-offload", action="store_true", + help="ZeRO-Inference: offload parameters to CPU.") + ds.add_argument("--nvme-offload-dir", default=None, + help="ZeRO-Inference: offload parameters to this NVMe directory.") + ds.add_argument("--local_rank", type=int, default=int(os.environ.get("LOCAL_RANK", 0)), + help="Set by the DeepSpeed launcher.") + + parser.add_argument("--output", required=True, help="Destination .jsonl file.") + return parser.parse_args() + + +def load_prompts(args, tokenizer) -> List[str]: + """Chat-templated prompts, ready to be fed to the model.""" + from rac.calibration import chat_template_prompt + + if args.dataset.endswith(".jsonl") or args.dataset.endswith(".json"): + raw = [] + with open(args.dataset) as handle: + for line in handle: + line = line.strip() + if line: + raw.append(json.loads(line)[args.prompt_column]) + else: + from datasets import load_dataset + + dataset = load_dataset(args.dataset, args.dataset_config, split=args.dataset_split) + if args.prompt_column not in dataset.column_names: + raise KeyError(f"Column '{args.prompt_column}' not in {args.dataset} " + f"(available: {dataset.column_names}). Set --prompt-column.") + raw = list(dataset[args.prompt_column]) + + if args.max_prompts is not None: + raw = raw[:args.max_prompts] + + return [ + chat_template_prompt(tokenizer, p, args.system_prompt, not args.no_chat_template) + for p in raw if p + ] + + +def generate_vllm(args, prompts: List[str], writer) -> int: + """Roll out with vLLM, in chunks, until the token budget is exhausted.""" + from vllm import LLM, SamplingParams + + llm = LLM( + model=args.model, + dtype=args.dtype, + tensor_parallel_size=args.tensor_parallel_size, + gpu_memory_utilization=args.gpu_memory_utilization, + trust_remote_code=args.trust_remote_code, + seed=args.seed, + ) + sampling = SamplingParams( + n=args.num_generations, + temperature=args.temperature, + top_p=args.top_p, + max_tokens=args.max_new_tokens, + seed=args.seed, + ) + + total = 0 + for start in range(0, len(prompts), args.batch_size): + chunk = prompts[start:start + args.batch_size] + for output in llm.generate(chunk, sampling): + prompt_tokens = len(output.prompt_token_ids) + for candidate in output.outputs: + total = writer(output.prompt, candidate.text, prompt_tokens, + len(candidate.token_ids)) + if total >= args.trace_tokens: + break + return total + + +def build_hf_model(args): + """Plain ``transformers`` model on a single device.""" + import torch + from transformers import AutoModelForCausalLM + + from rac import dtype_kwargs + + model = AutoModelForCausalLM.from_pretrained( + args.model, + trust_remote_code=args.trust_remote_code, + **dtype_kwargs(getattr(torch, args.dtype)), + ) + model.eval() + return model.to("cuda" if torch.cuda.is_available() else "cpu") + + +def build_deepspeed_model(args): + """ZeRO-Inference: ZeRO-3 sharding with optional CPU/NVMe parameter offload. + + Follows ``inference/huggingface/zero_inference/run_model.py`` in this repo. + All ranks execute the same ``generate`` call with ``synced_gpus=True``; + only rank 0 writes the traces. + """ + import deepspeed + import torch + from transformers import AutoConfig, AutoModelForCausalLM + from transformers.integrations.deepspeed import HfDeepSpeedConfig + + deepspeed.init_distributed() + config = AutoConfig.from_pretrained(args.model, trust_remote_code=args.trust_remote_code) + hidden_size = config.hidden_size + dtype = getattr(torch, args.dtype) + + ds_config = { + "fp16": {"enabled": dtype == torch.float16}, + "bf16": {"enabled": dtype == torch.bfloat16}, + "zero_optimization": { + "stage": 3, + "stage3_prefetch_bucket_size": 2 * hidden_size * hidden_size, + "stage3_param_persistence_threshold": hidden_size, + "stage3_max_live_parameters": 2 * hidden_size * hidden_size, + }, + "train_batch_size": args.batch_size * int(os.environ.get("WORLD_SIZE", 1)), + "steps_per_print": 2000, + "wall_clock_breakdown": False, + } + + if args.nvme_offload_dir: + Path(args.nvme_offload_dir).mkdir(parents=True, exist_ok=True) + ds_config["zero_optimization"]["offload_param"] = { + "device": "nvme", + "nvme_path": args.nvme_offload_dir, + "pin_memory": True, + "buffer_count": 5, + "buffer_size": 2 * (1024**3), + } + ds_config["aio"] = { + "block_size": 1048576 * 16, + "queue_depth": 64, + "thread_count": 8, + "single_submit": False, + "overlap_events": True, + } + elif args.cpu_offload: + ds_config["zero_optimization"]["offload_param"] = {"device": "cpu", "pin_memory": True} + + # Instructs from_pretrained to materialise the weights directly as ZeRO-3 + # shards; must stay alive until the model is built. + dschf = HfDeepSpeedConfig(ds_config) # noqa: F841 + + from rac import dtype_kwargs + + model = AutoModelForCausalLM.from_pretrained( + args.model, + trust_remote_code=args.trust_remote_code, + **dtype_kwargs(dtype), + ) + model.eval() + + engine = deepspeed.initialize(model=model, config_params=ds_config)[0] + engine.module.eval() + return engine.module + + +def generate_transformers(args, prompts: List[str], writer, use_deepspeed: bool) -> int: + """Roll out with ``model.generate`` (HF or ZeRO-Inference backend).""" + import torch + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(args.model, + trust_remote_code=args.trust_remote_code, + padding_side="left") + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + + model = build_deepspeed_model(args) if use_deepspeed else build_hf_model(args) + device = next(model.parameters()).device if not use_deepspeed else torch.device( + "cuda", args.local_rank) + is_writer_rank = int(os.environ.get("RANK", 0)) == 0 + + # Every rank counts tokens, but only rank 0 writes: under ZeRO-Inference all + # ranks must run the same number of generate() calls or they deadlock, so + # the stopping condition cannot depend on who is writing. + total = 0 + torch.manual_seed(args.seed) + for start in range(0, len(prompts), args.batch_size): + chunk = prompts[start:start + args.batch_size] + encoded = tokenizer(chunk, return_tensors="pt", padding=True, add_special_tokens=False) + # Some tokenizers also return token_type_ids, which causal LMs reject. + batch = { + k: v.to(device) + for k, v in encoded.items() if k in ("input_ids", "attention_mask") + } + + with torch.no_grad(): + generated = model.generate( + **batch, + do_sample=args.temperature > 0, + temperature=args.temperature, + top_p=args.top_p, + max_new_tokens=args.max_new_tokens, + num_return_sequences=args.num_generations, + pad_token_id=tokenizer.pad_token_id, + synced_gpus=use_deepspeed, + ) + + prompt_len = batch["input_ids"].shape[1] + completions = generated[:, prompt_len:] + for i, sequence in enumerate(completions): + prompt = chunk[i // args.num_generations] + text = tokenizer.decode(sequence, skip_special_tokens=True) + n_completion = int((sequence != tokenizer.pad_token_id).sum()) + n_prompt = int(batch["attention_mask"][i // args.num_generations].sum()) + total += n_prompt + n_completion + if is_writer_rank: + writer(prompt, text, n_prompt, n_completion) + + if total >= args.trace_tokens: + break + return total + + +def main() -> None: + args = parse_args() + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(args.model, + trust_remote_code=args.trust_remote_code) + prompts = load_prompts(args, tokenizer) + print(f"[traces] {len(prompts)} prompts loaded from {args.dataset}") + + backend = args.backend + if backend == "auto": + backend = "vllm" if importlib.util.find_spec("vllm") else "hf" + print(f"[traces] backend auto-selected: {backend}") + + output = Path(args.output) + output.parent.mkdir(parents=True, exist_ok=True) + + written = {"traces": 0, "tokens": 0} + started = time.time() + + with output.open("w") as handle: + + def writer(prompt: str, completion: str, n_prompt: int, n_completion: int) -> int: + """Append one trace; returns the running token total.""" + handle.write( + json.dumps({ + "prompt": prompt, + "completion": completion, + "prompt_tokens": n_prompt, + "completion_tokens": n_completion, + "model": args.model, + "dataset": args.dataset, + }) + "\n") + handle.flush() + written["traces"] += 1 + written["tokens"] += n_prompt + n_completion + if written["traces"] % 10 == 0: + print(f"[traces] {written['traces']} traces, " + f"{written['tokens']:,}/{args.trace_tokens:,} tokens, " + f"{time.time() - started:.0f}s") + return written["tokens"] + + if backend == "vllm": + generate_vllm(args, prompts, writer) + else: + generate_transformers(args, prompts, writer, use_deepspeed=backend == "deepspeed") + + print(f"[traces] wrote {written['traces']} traces " + f"({written['tokens']:,} tokens) to {output}") + if written["tokens"] < args.trace_tokens: + print(f"[traces] WARNING: token budget {args.trace_tokens:,} not reached; " + "raise --max-prompts or use a larger prompt dataset.") + + +if __name__ == "__main__": + main() diff --git a/compression/reasoning_aware_compression/prune.py b/compression/reasoning_aware_compression/prune.py new file mode 100644 index 000000000..64afc622c --- /dev/null +++ b/compression/reasoning_aware_compression/prune.py @@ -0,0 +1,218 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""Phase II of RAC: one-shot pruning against a chosen calibration set. + +Everything except ``--calibration`` is a standard SparseGPT/Wanda run; passing +``--calibration rac --traces `` is what turns it into Reasoning-Aware +Compression. + +Example -- 50% unstructured sparsity with RAC calibration:: + + python prune.py \\ + --model deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B \\ + --calibration rac --traces traces/math_1.5b.jsonl \\ + --sparsity 0.5 --nsamples 512 --seqlen 2048 \\ + --output checkpoints/r1-qwen-1.5b-rac-50 + +The model is held on CPU and streamed to the accelerator one transformer block +at a time, so peak device memory is set by the largest block rather than by the +model. +""" + +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path + +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer + +from rac import ( + build_calibration_samples, + dtype_kwargs, + layerwise_sparsity, + magnitude_prune, + model_sparsity, + select_blocks_by_thirds, + sequential_prune, +) +from rac.calibration import CALIBRATION_SOURCES + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + + model = parser.add_argument_group("model") + model.add_argument("--model", required=True, help="Dense checkpoint to prune.") + model.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "float16", "float32"]) + model.add_argument("--trust-remote-code", action="store_true") + model.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") + + calib = parser.add_argument_group("calibration") + calib.add_argument("--calibration", default="rac", choices=list(CALIBRATION_SOURCES), + help="rac: prompts + on-policy CoT traces (the method). " + "prompt: prompts only. c4: generic web text.") + calib.add_argument("--traces", default=None, + help="JSONL from collect_traces.py. Required for --calibration rac; " + "with --calibration prompt it reuses the same prompts without CoT.") + calib.add_argument("--dataset", default=None, + help="HF dataset (or .jsonl) for --calibration prompt / c4.") + calib.add_argument("--dataset-config", default=None) + calib.add_argument("--dataset-split", default="train") + calib.add_argument("--prompt-column", default="problem") + calib.add_argument("--system-prompt", default=None) + calib.add_argument("--no-chat-template", action="store_true") + calib.add_argument("--nsamples", type=int, default=512, + help="Calibration windows. nsamples * seqlen = calibration tokens; " + "512 x 2048 = 1M, as in the paper.") + calib.add_argument("--seqlen", type=int, default=2048, help="Tokens per calibration window.") + calib.add_argument("--seed", type=int, default=0) + + prune = parser.add_argument_group("pruning") + prune.add_argument("--pruning-method", default="sparsegpt", + choices=["sparsegpt", "wanda", "magnitude"]) + prune.add_argument("--sparsity", type=float, default=0.5, + help="Per-layer unstructured sparsity; ignored with --prune-n/--prune-m.") + prune.add_argument("--prune-n", type=int, default=0, + help="Semi-structured n:m sparsity, e.g. --prune-n 2 --prune-m 4.") + prune.add_argument("--prune-m", type=int, default=0) + prune.add_argument("--scope", default="all", choices=["all", "mlp"], + help="Which projections to prune within a block.") + prune.add_argument("--layer-thirds", default="1,2,3", + help="Which thirds of the decoder stack to prune, e.g. '1,3' for the " + "first and last third (Table 6 of the paper).") + prune.add_argument("--batch-size", type=int, default=1, + help="Calibration samples per forward pass.") + prune.add_argument("--blocksize", type=int, default=128, help="SparseGPT column block size.") + prune.add_argument("--percdamp", type=float, default=0.01, + help="Hessian damping, as a fraction of its mean diagonal.") + + parser.add_argument("--output", required=True, help="Directory for the pruned checkpoint.") + parser.add_argument("--save-report", action="store_true", + help="Also write per-layer sparsity to the output directory.") + return parser.parse_args() + + +def validate(args: argparse.Namespace) -> None: + if bool(args.prune_n) != bool(args.prune_m): + raise ValueError("--prune-n and --prune-m must be given together (e.g. 2 and 4).") + if args.prune_n and args.prune_n >= args.prune_m: + raise ValueError(f"--prune-n ({args.prune_n}) must be smaller than " + f"--prune-m ({args.prune_m}).") + if not args.prune_n and not 0.0 < args.sparsity < 1.0: + raise ValueError(f"--sparsity must be in (0, 1), got {args.sparsity}.") + if args.calibration == "rac" and not args.traces: + raise ValueError("--calibration rac requires --traces; run collect_traces.py first.") + + +def main() -> None: + args = parse_args() + validate(args) + + print(f"[rac] loading {args.model}") + tokenizer = AutoTokenizer.from_pretrained(args.model, + trust_remote_code=args.trust_remote_code) + model = AutoModelForCausalLM.from_pretrained( + args.model, + trust_remote_code=args.trust_remote_code, + low_cpu_mem_usage=True, + **dtype_kwargs(getattr(torch, args.dtype)), + ) + model.eval() + for param in model.parameters(): + param.requires_grad_(False) + + thirds = [int(t) for t in args.layer_thirds.replace(",", " ").split()] + num_blocks = model.config.num_hidden_layers + block_indices = select_blocks_by_thirds(num_blocks, thirds) + target = (f"{args.prune_n}:{args.prune_m}" if args.prune_n else f"{args.sparsity:.0%}") + print(f"[rac] pruning {len(block_indices)}/{num_blocks} blocks to {target} " + f"with {args.pruning_method}, scope={args.scope}") + + started = time.time() + stats = [] + + if args.pruning_method == "magnitude": + print("[rac] magnitude pruning ignores calibration data") + magnitude_prune( + model, + sparsity=args.sparsity, + prune_n=args.prune_n, + prune_m=args.prune_m, + scope=args.scope, + block_indices=block_indices, + ) + else: + print(f"[rac] building '{args.calibration}' calibration set: " + f"{args.nsamples} x {args.seqlen} = {args.nsamples * args.seqlen:,} tokens") + samples = build_calibration_samples( + args.calibration, + tokenizer, + nsamples=args.nsamples, + seqlen=args.seqlen, + seed=args.seed, + traces=args.traces, + dataset_name=args.dataset, + dataset_config=args.dataset_config, + dataset_split=args.dataset_split, + prompt_column=args.prompt_column, + system_prompt=args.system_prompt, + use_chat_template=not args.no_chat_template, + ) + stats = sequential_prune( + model, + samples, + method=args.pruning_method, + sparsity=args.sparsity, + prune_n=args.prune_n, + prune_m=args.prune_m, + scope=args.scope, + block_indices=block_indices, + device=args.device, + blocksize=args.blocksize, + percdamp=args.percdamp, + batch_size=args.batch_size, + ) + + elapsed = time.time() - started + realised = model_sparsity(model, scope=args.scope, block_indices=block_indices) + overall = model_sparsity(model, scope="all") + print(f"[rac] done in {elapsed / 60:.1f} min; realised sparsity {realised:.2%} over the " + f"pruned layers, {overall:.2%} over all decoder-block weights") + + output = Path(args.output) + output.mkdir(parents=True, exist_ok=True) + model.save_pretrained(output) + tokenizer.save_pretrained(output) + + metadata = { + "base_model": args.model, + "pruning_method": args.pruning_method, + "calibration": args.calibration, + "traces": args.traces, + "calibration_tokens": args.nsamples * args.seqlen, + "sparsity": args.sparsity, + "prune_n": args.prune_n, + "prune_m": args.prune_m, + "scope": args.scope, + "layer_thirds": thirds, + "pruned_blocks": block_indices, + "realised_sparsity_pruned_layers": realised, + "realised_sparsity_all_blocks": overall, + "wall_clock_seconds": elapsed, + "reconstruction_loss": sum(s["loss"] for s in stats) if stats else None, + } + (output / "rac_pruning.json").write_text(json.dumps(metadata, indent=2) + "\n") + + if args.save_report: + report = dict(layerwise_sparsity(model)) + (output / "rac_layer_sparsity.json").write_text(json.dumps(report, indent=2) + "\n") + + print(f"[rac] pruned checkpoint written to {output}") + + +if __name__ == "__main__": + main() diff --git a/compression/reasoning_aware_compression/rac/__init__.py b/compression/reasoning_aware_compression/rac/__init__.py new file mode 100644 index 000000000..003b90c41 --- /dev/null +++ b/compression/reasoning_aware_compression/rac/__init__.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""Reasoning-Aware Compression (RAC): one-shot pruning of reasoning LLMs. + +Implements `Reasoning Models Can be Accurately Pruned Via Chain-of-Thought +Reconstruction `_ (Lucas et al., ICLR 2026). +""" + +from .calibration import CALIBRATION_SOURCES, build_calibration_samples +from .magnitude import magnitude_prune +from .modelutils import ( + dtype_kwargs, + find_linear_layers, + get_decoder_layers, + layerwise_sparsity, + model_sparsity, + prunable_layers, + select_blocks_by_thirds, +) +from .sequential import sequential_prune +from .sparsegpt import SparseGPT +from .wanda import Wanda + +__all__ = [ + "CALIBRATION_SOURCES", + "SparseGPT", + "Wanda", + "build_calibration_samples", + "dtype_kwargs", + "find_linear_layers", + "get_decoder_layers", + "layerwise_sparsity", + "magnitude_prune", + "model_sparsity", + "prunable_layers", + "select_blocks_by_thirds", + "sequential_prune", +] diff --git a/compression/reasoning_aware_compression/rac/calibration.py b/compression/reasoning_aware_compression/rac/calibration.py new file mode 100644 index 000000000..fd07734a2 --- /dev/null +++ b/compression/reasoning_aware_compression/rac/calibration.py @@ -0,0 +1,257 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""Calibration sets for one-shot pruning -- this is where RAC happens. + +Layer-wise pruning minimises ``||W X - W' X||_F^2`` over a calibration +activation matrix ``X``. Three ways of building the token stream behind ``X`` +are supported, matching the ablation in the RAC paper: + +``c4`` + Generic web text (the SparseGPT/Wanda default). Nothing about the task or + the model's own behaviour is represented. +``prompt`` + Task prompts only (e.g. OpenR1-Math-220k problems), chat-templated the way + the model will see them at inference. Covers ``X^P`` from the paper. +``rac`` + Task prompts *concatenated with the dense model's own on-policy + chain-of-thought*, collected by ``collect_traces.py``. Covers + ``X^RAC = [X^P, X^D]``. + +Teacher-forcing a sampled trace reproduces exactly the hidden states the model +computed while generating it, so packing traces into the calibration stream +recovers the decode-time activations of Algorithm 1 without re-running +generation during pruning. +""" + +from __future__ import annotations + +import json +import random +from pathlib import Path +from typing import Iterable, Iterator, List, Optional, Sequence + +import torch + +CALIBRATION_SOURCES = ("c4", "prompt", "rac") + +# Enough documents to fill a 1M-token budget with room to spare, so that the +# shuffle in :func:`pack_token_windows` has something to choose from. +_DEFAULT_TEXT_LIMIT = 20000 + + +def chat_template_prompt( + tokenizer, + prompt: str, + system_prompt: Optional[str] = None, + use_chat_template: bool = True, +) -> str: + """Render a raw prompt the way the model will receive it at inference.""" + if not use_chat_template or tokenizer.chat_template is None: + return prompt + + messages = [] + if system_prompt: + messages.append({"role": "system", "content": system_prompt}) + messages.append({"role": "user", "content": prompt}) + return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + + +def load_trace_texts(path: str, with_completion: bool = True) -> Iterator[str]: + """Read the JSONL written by ``collect_traces.py``. + + Each line holds the chat-templated ``prompt`` and the model's + ``completion``. ``with_completion=False`` yields the prompt-only variant of + the *same* prompts, which is the apples-to-apples baseline for RAC. + """ + trace_path = Path(path) + if not trace_path.exists(): + raise FileNotFoundError(f"Trace file '{path}' not found. Run collect_traces.py first.") + + with trace_path.open() as handle: + for line in handle: + line = line.strip() + if not line: + continue + row = json.loads(line) + text = row["prompt"] + if with_completion: + text = text + row["completion"] + if text: + yield text + + +def load_prompt_texts( + dataset_name: str, + prompt_column: str, + tokenizer, + split: str = "train", + dataset_config: Optional[str] = None, + system_prompt: Optional[str] = None, + use_chat_template: bool = True, + limit: int = _DEFAULT_TEXT_LIMIT, +) -> Iterator[str]: + """Chat-templated task prompts from a Hugging Face dataset or a JSONL file.""" + for i, prompt in enumerate(_iter_raw_prompts(dataset_name, prompt_column, split, + dataset_config)): + if i >= limit: + break + if prompt: + yield chat_template_prompt(tokenizer, prompt, system_prompt, use_chat_template) + + +def load_c4_texts( + dataset_name: str = "allenai/c4", + dataset_config: str = "en", + split: str = "train", + limit: int = _DEFAULT_TEXT_LIMIT, +) -> Iterator[str]: + """Stream raw web text -- the calibration set RAC is measured against.""" + from datasets import load_dataset + + stream = load_dataset(dataset_name, dataset_config, split=split, streaming=True) + for i, row in enumerate(stream): + if i >= limit: + break + text = row.get("text") + if text: + yield text + + +def pack_token_windows( + texts: Iterable[str], + tokenizer, + nsamples: int, + seqlen: int, + seed: int = 0, + shuffle_buffer: int = 4096, +) -> List[torch.Tensor]: + """Tokenise ``texts`` and cut the stream into ``nsamples`` windows. + + Documents are separated by EOS and concatenated, then chunked into + fixed-length windows. Fixed length keeps the causal mask and rotary + embeddings identical across samples, which is what lets + :mod:`rac.sequential` capture the block kwargs once, and it avoids padding + tokens polluting the Hessian. + + Reasoning traces are typically several thousand tokens, so in ``rac`` mode + most windows come from a single trace. + """ + if nsamples <= 0 or seqlen <= 0: + raise ValueError("nsamples and seqlen must both be positive.") + + rng = random.Random(seed) + buffer: List[str] = [] + samples: List[torch.Tensor] = [] + stream: List[int] = [] + eos_id = tokenizer.eos_token_id + + def flush(docs: Sequence[str]) -> None: + for doc in docs: + ids = tokenizer(doc, add_special_tokens=False).input_ids + stream.extend(ids) + if eos_id is not None: + stream.append(eos_id) + while len(stream) >= seqlen and len(samples) < nsamples: + samples.append(torch.tensor(stream[:seqlen], dtype=torch.long)) + del stream[:seqlen] + if len(samples) >= nsamples: + return + + # Shuffle within a bounded buffer so that a small token budget still draws + # from across the corpus without materialising all of it. + for text in texts: + buffer.append(text) + if len(buffer) >= shuffle_buffer: + rng.shuffle(buffer) + flush(buffer) + buffer = [] + if len(samples) >= nsamples: + break + + if len(samples) < nsamples and buffer: + rng.shuffle(buffer) + flush(buffer) + + if len(samples) < nsamples: + raise ValueError(f"Calibration corpus yielded only {len(samples)} windows of {seqlen} " + f"tokens, {nsamples} requested. Collect more traces, lower --nsamples, " + "or lower --seqlen.") + return samples + + +def build_calibration_samples( + source: str, + tokenizer, + nsamples: int, + seqlen: int, + seed: int = 0, + traces: Optional[str] = None, + dataset_name: Optional[str] = None, + dataset_config: Optional[str] = None, + dataset_split: str = "train", + prompt_column: str = "problem", + system_prompt: Optional[str] = None, + use_chat_template: bool = True, +) -> List[torch.Tensor]: + """Build ``nsamples`` calibration windows of ``seqlen`` tokens. + + ``source="prompt"`` reads prompts from ``traces`` when given (so the + prompt-only baseline sees exactly the prompts RAC saw) and falls back to + ``dataset_name`` otherwise. + """ + if source not in CALIBRATION_SOURCES: + raise ValueError(f"Unknown calibration source '{source}'; " + f"expected one of {list(CALIBRATION_SOURCES)}.") + + if source == "rac": + if not traces: + raise ValueError("--calibration rac requires --traces; run collect_traces.py first.") + texts: Iterable[str] = load_trace_texts(traces, with_completion=True) + elif source == "prompt": + if traces: + texts = load_trace_texts(traces, with_completion=False) + elif dataset_name: + texts = load_prompt_texts( + dataset_name, + prompt_column, + tokenizer, + split=dataset_split, + dataset_config=dataset_config, + system_prompt=system_prompt, + use_chat_template=use_chat_template, + ) + else: + raise ValueError("--calibration prompt requires either --traces or --dataset.") + else: + texts = load_c4_texts( + dataset_name or "allenai/c4", + dataset_config or "en", + dataset_split, + ) + + return pack_token_windows(texts, tokenizer, nsamples, seqlen, seed=seed) + + +def _iter_raw_prompts( + dataset_name: str, + prompt_column: str, + split: str, + dataset_config: Optional[str], +) -> Iterator[str]: + """Yield the raw prompt strings of a HF dataset or a local JSONL file.""" + if dataset_name.endswith(".jsonl") or dataset_name.endswith(".json"): + with open(dataset_name) as handle: + for line in handle: + line = line.strip() + if line: + yield json.loads(line)[prompt_column] + return + + from datasets import load_dataset + + dataset = load_dataset(dataset_name, dataset_config, split=split) + if prompt_column not in dataset.column_names: + raise KeyError(f"Column '{prompt_column}' not in {dataset_name} " + f"(available: {dataset.column_names}). Set --prompt-column.") + for row in dataset: + yield row[prompt_column] diff --git a/compression/reasoning_aware_compression/rac/magnitude.py b/compression/reasoning_aware_compression/rac/magnitude.py new file mode 100644 index 000000000..c3ff50532 --- /dev/null +++ b/compression/reasoning_aware_compression/rac/magnitude.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""Calibration-free layer-wise magnitude pruning. + +The weakest baseline in the RAC paper (``MP``). It looks only at weight +magnitudes, so the calibration set -- and therefore RAC -- has no effect on it. +Useful as a floor when reporting numbers. +""" + +from __future__ import annotations + +from typing import Optional, Sequence + +import torch +import torch.nn as nn + +from .modelutils import find_linear_layers, get_decoder_layers + + +@torch.no_grad() +def magnitude_prune( + model: nn.Module, + sparsity: float, + prune_n: int = 0, + prune_m: int = 0, + scope: str = "all", + block_indices: Optional[Sequence[int]] = None, +) -> None: + """Prune ``model`` in place, independently per layer.""" + blocks = get_decoder_layers(model) + selected = range(len(blocks)) if block_indices is None else block_indices + + for i in selected: + for _, layer in find_linear_layers(blocks[i], scope=scope).items(): + W = layer.weight.data + metric = W.abs().float() + + if prune_n != 0: + mask = torch.zeros_like(metric, dtype=torch.bool) + for col in range(0, W.shape[1], prune_m): + group = metric[:, col:col + prune_m] + idx = torch.topk(group, prune_n, dim=1, largest=False)[1] + mask[:, col:col + prune_m].scatter_(1, idx, True) + else: + k = int(W.numel() * sparsity) + if k == 0: + continue + thresh = torch.kthvalue(metric.flatten(), k).values + mask = metric <= thresh + + W[mask] = 0 + print(f"[magnitude] block {i} pruned") diff --git a/compression/reasoning_aware_compression/rac/modelutils.py b/compression/reasoning_aware_compression/rac/modelutils.py new file mode 100644 index 000000000..8aab90277 --- /dev/null +++ b/compression/reasoning_aware_compression/rac/modelutils.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""Small model-introspection helpers shared by the RAC pruners.""" + +from __future__ import annotations + +from typing import Dict, Iterable, List, Optional, Tuple + +import torch +import torch.nn as nn + +# Sub-modules we are willing to prune. Attention/MLP projections in every +# decoder-only architecture we target (Qwen2/Qwen3, Llama, Mistral) are plain +# ``nn.Linear``. +PRUNABLE_TYPES = (nn.Linear, ) + +MLP_KEYWORDS = ("mlp", "feed_forward", "ffn", "gate_proj", "up_proj", "down_proj") + + +def dtype_kwargs(dtype: torch.dtype) -> Dict[str, torch.dtype]: + """``from_pretrained`` dtype argument, named for the installed transformers. + + The argument was renamed from ``torch_dtype`` to ``dtype`` in transformers + 4.56; the old name still works but warns, and is slated for removal. + """ + from packaging.version import parse + from transformers import __version__ as transformers_version + + if parse(transformers_version).release >= (4, 56): + return {"dtype": dtype} + return {"torch_dtype": dtype} + + +def get_decoder_layers(model: nn.Module) -> nn.ModuleList: + """Return the stack of transformer blocks, independent of model layout. + + Handles ``model.model.layers`` (Llama/Qwen/Mistral) and + ``model.model.decoder.layers`` (OPT). + """ + core = getattr(model, "model", model) + core = getattr(core, "decoder", core) + layers = getattr(core, "layers", None) + if layers is None: + raise AttributeError(f"Cannot locate decoder layers on {type(model).__name__}; " + "RAC supports decoder-only causal LMs.") + return layers + + +def get_embedding_modules(model: nn.Module) -> List[nn.Module]: + """Modules that run *before* the first decoder block. + + These have to sit on the calibration device while we capture the inputs of + block 0, and can go back to CPU afterwards. + """ + core = getattr(model, "model", model) + core = getattr(core, "decoder", core) + names = ("embed_tokens", "embed_positions", "rotary_emb", "embed_layer_norm") + return [getattr(core, n) for n in names if getattr(core, n, None) is not None] + + +def find_linear_layers(module: nn.Module, scope: str = "all") -> Dict[str, nn.Module]: + """Collect prunable sub-modules of ``module`` keyed by qualified name. + + ``scope="mlp"`` restricts the result to feed-forward projections, which is + the setting used for the 2:4 throughput experiments in the RAC paper (the + attention projections are left dense). + """ + if scope not in ("all", "mlp"): + raise ValueError(f"Unknown scope '{scope}', expected 'all' or 'mlp'.") + + found = {} + for name, sub in module.named_modules(): + if not isinstance(sub, PRUNABLE_TYPES): + continue + if scope == "mlp" and not any(k in name.lower() for k in MLP_KEYWORDS): + continue + found[name] = sub + return found + + +def select_blocks_by_thirds(num_blocks: int, thirds: Iterable[int]) -> List[int]: + """Indices of the decoder blocks that fall in the requested thirds. + + Thirds are 1-based: ``(1,)`` is the first third of the network, ``(1, 3)`` + the first and last thirds, ``(1, 2, 3)`` the whole model. + """ + thirds = sorted({int(t) for t in thirds}) + for t in thirds: + if t not in (1, 2, 3): + raise ValueError(f"Invalid third '{t}', expected values in {{1, 2, 3}}.") + + a, b = num_blocks // 3, (2 * num_blocks) // 3 + bounds = {1: (0, a), 2: (a, b), 3: (b, num_blocks)} + selected: List[int] = [] + for t in thirds: + lo, hi = bounds[t] + selected.extend(range(lo, hi)) + return selected + + +def prunable_layers(model: nn.Module, + scope: str = "all", + block_indices: Optional[Iterable[int]] = None) -> Dict[str, nn.Module]: + """Prunable layers of the whole model, keyed by fully-qualified name. + + Only layers *inside decoder blocks* are returned: the embedding matrix and + the LM head are never pruned, so counting them would misreport sparsity. + """ + blocks = get_decoder_layers(model) + selected = range(len(blocks)) if block_indices is None else block_indices + + wanted = set() + for i in selected: + wanted.update(id(layer) for layer in find_linear_layers(blocks[i], scope=scope).values()) + + return { + name: module + for name, module in model.named_modules() + if isinstance(module, PRUNABLE_TYPES) and id(module) in wanted + } + + +@torch.no_grad() +def layerwise_sparsity(model: nn.Module, + scope: str = "all", + block_indices: Optional[Iterable[int]] = None) -> List[Tuple[str, float]]: + """Fraction of exactly-zero weights in every prunable layer.""" + report = [] + for name, layer in prunable_layers(model, scope=scope, block_indices=block_indices).items(): + w = layer.weight.data + report.append((name, (w == 0).sum().item() / w.numel())) + return report + + +@torch.no_grad() +def model_sparsity(model: nn.Module, + scope: str = "all", + block_indices: Optional[Iterable[int]] = None) -> float: + """Fraction of exactly-zero weights across the selected prunable layers.""" + zeros, total = 0, 0 + for layer in prunable_layers(model, scope=scope, block_indices=block_indices).values(): + w = layer.weight.data + zeros += (w == 0).sum().item() + total += w.numel() + return zeros / total if total else 0.0 diff --git a/compression/reasoning_aware_compression/rac/sequential.py b/compression/reasoning_aware_compression/rac/sequential.py new file mode 100644 index 000000000..10abf3a52 --- /dev/null +++ b/compression/reasoning_aware_compression/rac/sequential.py @@ -0,0 +1,266 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""Block-by-block calibration driver for one-shot pruning. + +Only one transformer block is resident on the accelerator at a time: the +calibration activations flow through the network block by block, and each +block is pruned before its (already sparse) outputs are propagated to the +next one. Peak device memory is therefore ``one block + its Hessians + +one activation slice``, which is what lets a 70B checkpoint be pruned on a +single 80GB GPU, as in the RAC paper. + +This module is deliberately agnostic to *what* the calibration tokens are -- +prompt-only, C4, or prompt + on-policy chain-of-thought (RAC). Swapping the +calibration set is the entire method; see :mod:`rac.calibration`. +""" + +from __future__ import annotations + +from typing import Dict, List, Optional, Sequence, Tuple + +import torch +import torch.nn as nn + +from .modelutils import find_linear_layers, get_decoder_layers, get_embedding_modules +from .sparsegpt import SparseGPT, hessian_bytes +from .wanda import Wanda + +# Keyword arguments a decoder block accepts that we want to replay for every +# calibration sample. Anything else (KV caches in particular) is dropped. +_REPLAYED_BLOCK_KWARGS = ( + "attention_mask", + "position_ids", + "position_embeddings", + "cache_position", +) + + +class _StopForward(Exception): + """Raised by the catcher to abort the forward pass after block 0's input.""" + + +class _Catcher(nn.Module): + """Records the hidden states and kwargs that reach the first block.""" + + def __init__(self, inner: nn.Module, store: Dict): + super().__init__() + self.inner = inner + self.store = store + + def forward(self, hidden_states, **kwargs): + idx = self.store["index"] + self.store["inps"][idx] = hidden_states.squeeze(0).to(self.store["inps"].device) + self.store["index"] = idx + 1 + if self.store["kwargs"] is None: + self.store["kwargs"] = { + k: v + for k, v in kwargs.items() if k in _REPLAYED_BLOCK_KWARGS and v is not None + } + raise _StopForward + + def __getattr__(self, name): + # nn.Module.__getattr__ only runs for attributes the module does not + # own, so this forwards lookups like ``attention_type`` that the + # surrounding model performs on the block it wraps. + try: + return super().__getattr__(name) + except AttributeError: + return getattr(self.inner, name) + + +def get_pruner_factory(method: str): + """Map a ``--pruning-method`` name to a per-layer pruner class.""" + factories = {"sparsegpt": SparseGPT, "wanda": Wanda} + if method not in factories: + raise ValueError(f"Unknown calibrated pruning method '{method}'; " + f"expected one of {sorted(factories)}.") + return factories[method] + + +@torch.no_grad() +def capture_block_inputs( + model: nn.Module, + samples: Sequence[torch.Tensor], + device: torch.device, +) -> Tuple[torch.Tensor, Dict]: + """Run the embedding stack and collect the inputs of decoder block 0. + + Returns ``(inps, block_kwargs)`` where ``inps`` is a CPU tensor of shape + ``(nsamples, seqlen, hidden)``. Because every calibration sample has the + same length, the block kwargs (causal mask, position ids, rotary + embeddings) are identical across samples and captured once. + """ + layers = get_decoder_layers(model) + hidden_size = model.config.hidden_size + dtype = next(model.parameters()).dtype + seqlen = samples[0].numel() + + store = { + "inps": torch.zeros((len(samples), seqlen, hidden_size), dtype=dtype, device="cpu"), + "kwargs": None, + "index": 0, + } + + embeddings = get_embedding_modules(model) + for module in embeddings: + module.to(device) + + layers[0] = _Catcher(layers[0], store) + try: + for sample in samples: + if sample.numel() != seqlen: + raise ValueError("All calibration samples must have the same length; " + f"got {sample.numel()} and {seqlen}.") + try: + model(sample.reshape(1, -1).to(device)) + except _StopForward: + pass + finally: + layers[0] = layers[0].inner + for module in embeddings: + module.to("cpu") + + if store["kwargs"] is None: + raise RuntimeError("No calibration sample reached the first decoder block.") + return store["inps"], store["kwargs"] + + +@torch.no_grad() +def sequential_prune( + model: nn.Module, + samples: Sequence[torch.Tensor], + method: str = "sparsegpt", + sparsity: float = 0.5, + prune_n: int = 0, + prune_m: int = 0, + scope: str = "all", + block_indices: Optional[Sequence[int]] = None, + device: str = "cuda", + blocksize: int = 128, + percdamp: float = 0.01, + batch_size: int = 1, +) -> List[dict]: + """Prune ``model`` in place, block by block, on ``samples``. + + Args: + samples: calibration sequences of identical length, as produced by + :func:`rac.calibration.build_calibration_samples`. + method: ``"sparsegpt"`` or ``"wanda"``. + sparsity: per-layer unstructured sparsity; ignored when ``prune_n > 0``. + prune_n, prune_m: semi-structured ``n:m`` pattern (e.g. 2:4). + scope: ``"all"`` prunes attention and MLP projections, ``"mlp"`` only + the MLP ones. + block_indices: decoder blocks to prune; ``None`` means all of them. + Blocks outside the list still propagate activations, they are just + left dense. + batch_size: calibration samples per forward pass. Higher is faster but + holds more activation memory on the device. + + Returns: + One dict per pruned layer with its block index, name and reconstruction + loss. + """ + device = torch.device(device) + pruner_cls = get_pruner_factory(method) + layers = get_decoder_layers(model) + selected = set(range(len(layers)) if block_indices is None else block_indices) + + use_cache = getattr(model.config, "use_cache", False) + model.config.use_cache = False + + inps, block_kwargs = capture_block_inputs(model, samples, device) + outs = torch.empty_like(inps) + stats: List[dict] = [] + + for i, block in enumerate(layers): + block.to(device) + + if i in selected: + targets = find_linear_layers(block, scope=scope) + else: + targets = {} + print(f"[rac] block {i}: left dense") + + if targets: + pruners = {name: pruner_cls(layer) for name, layer in targets.items()} + if method == "sparsegpt": + budget = sum(hessian_bytes(layer) for layer in targets.values()) + print(f"[rac] block {i}: {len(targets)} layers, " + f"{budget / 2**30:.2f} GiB of Hessians") + + # Pre-hooks see exactly the tensor each linear layer consumes, so + # this works regardless of how the block routes its activations. + handles = [ + layer.register_forward_pre_hook( + lambda _module, args, name=name: pruners[name].add_batch(args[0])) + for name, layer in targets.items() + ] + _run_block(block, inps, outs, block_kwargs, device, batch_size, store_out=False) + for handle in handles: + handle.remove() + + for name in targets: + loss = pruners[name].prune( + sparsity, + prune_n=prune_n, + prune_m=prune_m, + blocksize=blocksize, + percdamp=percdamp, + ) + pruners[name].free() + stats.append({"block": i, "layer": name, "loss": loss}) + print(f"[rac] block {i} :: {name} pruned (loss {loss:.4f})") + del pruners + + # Propagate the (now sparse) block outputs to the next block, so every + # layer is reconstructed against the activations it will really see. + _run_block(block, inps, outs, block_kwargs, device, batch_size, store_out=True) + + block.to("cpu") + _empty_cache(device) + inps, outs = outs, inps + + model.config.use_cache = use_cache + return stats + + +def _run_block( + block: nn.Module, + inps: torch.Tensor, + outs: torch.Tensor, + block_kwargs: Dict, + device: torch.device, + batch_size: int, + store_out: bool, +) -> None: + """Push every calibration sample through one block.""" + for start in range(0, inps.shape[0], batch_size): + end = min(start + batch_size, inps.shape[0]) + hidden = inps[start:end].to(device) + out = block(hidden, **_expand_kwargs(block_kwargs, end - start)) + if isinstance(out, tuple): + out = out[0] + if store_out: + outs[start:end] = out.to(outs.device, dtype=outs.dtype) + del hidden, out + + +def _expand_kwargs(block_kwargs: Dict, batch: int) -> Dict: + """Repeat the captured single-sample kwargs along the batch dimension.""" + if batch == 1: + return dict(block_kwargs) + + expanded = {} + for key, value in block_kwargs.items(): + if key == "position_embeddings" and isinstance(value, tuple): + expanded[key] = tuple(v.expand(batch, *v.shape[1:]) for v in value) + elif torch.is_tensor(value) and value.dim() >= 2 and value.shape[0] == 1: + expanded[key] = value.expand(batch, *value.shape[1:]) + else: + expanded[key] = value + return expanded + + +def _empty_cache(device: torch.device) -> None: + if device.type == "cuda" and torch.cuda.is_available(): + torch.cuda.empty_cache() diff --git a/compression/reasoning_aware_compression/rac/sparsegpt.py b/compression/reasoning_aware_compression/rac/sparsegpt.py new file mode 100644 index 000000000..62e370ff8 --- /dev/null +++ b/compression/reasoning_aware_compression/rac/sparsegpt.py @@ -0,0 +1,168 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""SparseGPT one-shot pruning of a single linear layer. + +Adapted from the reference implementation released with +`SparseGPT: Massive Language Models Can Be Accurately Pruned in One-Shot +`_ (Frantar & Alistarh, ICML 2023), +https://github.com/IST-DASLab/sparsegpt (Apache-2.0), and from the RAC +reference code at https://github.com/RyanLucas3/Reasoning-Aware-Compression. + +The algorithm itself is untouched by RAC: given the layer input activations +:math:`X_\\ell` it minimises :math:`\\|W_\\ell X_\\ell - \\widehat{W}_\\ell X_\\ell\\|_F^2` +subject to a sparsity constraint. What RAC changes is *which* activations end +up in :math:`X_\\ell` -- see :mod:`rac.calibration`. +""" + +from __future__ import annotations + +import math + +import torch +import torch.nn as nn + +# The Hessian accumulation and the Cholesky solve are numerically delicate; +# TF32 silently drops ~10 bits of mantissa and makes the factorisation fail on +# ill-conditioned layers. +torch.backends.cuda.matmul.allow_tf32 = False +torch.backends.cudnn.allow_tf32 = False + + +class SparseGPT: + """Accumulates :math:`H = 2/N \\cdot X X^\\top` for one layer, then prunes it.""" + + def __init__(self, layer: nn.Linear): + self.layer = layer + self.dev = layer.weight.device + self.rows, self.columns = layer.weight.data.shape + self.H = torch.zeros((self.columns, self.columns), device=self.dev, dtype=torch.float32) + self.nsamples = 0 + + @torch.no_grad() + def add_batch(self, inp: torch.Tensor) -> None: + """Fold one batch of layer inputs into the running Hessian. + + ``inp`` is ``(batch, seq, in_features)`` or ``(tokens, in_features)``; + every token contributes one column of :math:`X_\\ell`. + """ + if inp.dim() == 3: + inp = inp.reshape(-1, inp.shape[-1]) + elif inp.dim() == 1: + inp = inp.unsqueeze(0) + inp = inp.t() # (in_features, tokens) + + ntokens = inp.shape[1] + self.H *= self.nsamples / (self.nsamples + ntokens) + self.nsamples += ntokens + inp = math.sqrt(2 / self.nsamples) * inp.float() + self.H += inp @ inp.t() + + @torch.no_grad() + def prune( + self, + sparsity: float, + prune_n: int = 0, + prune_m: int = 0, + blocksize: int = 128, + percdamp: float = 0.01, + max_percdamp: float = 0.5, + ) -> float: + """Prune the layer in place and return the reconstruction loss. + + With ``prune_n``/``prune_m`` set (e.g. 2:4) the layer is made + semi-structured sparse and ``sparsity`` is ignored. + + The Hessian can be singular when a calibration set activates only part + of the input space, so the damping factor is doubled and the whole + solve retried whenever the Cholesky factorisation fails. + """ + if self.nsamples == 0: + raise RuntimeError("SparseGPT.prune() called before any calibration batch.") + + w_base = self.layer.weight.data.clone() + damp_scale = percdamp + + while True: + H = self.H.clone() + W = w_base.clone().float() + + # Columns that never fire cannot be reconstructed; zero them out + # and keep the Hessian invertible. + dead = torch.diag(H) == 0 + H[dead, dead] = 1 + W[:, dead] = 0 + + diag = torch.arange(self.columns, device=self.dev) + H[diag, diag] += damp_scale * torch.mean(torch.diag(H)) + + try: + L = torch.linalg.cholesky(H) + Hinv = torch.cholesky_inverse(L) + Hinv = torch.linalg.cholesky(Hinv, upper=True) + del L, H + except torch.linalg.LinAlgError as err: + if damp_scale >= max_percdamp: + raise RuntimeError( + f"SparseGPT: Cholesky failed up to percdamp={damp_scale:.3f}. " + "The calibration set is likely too small for this layer.") from err + damp_scale = min(damp_scale * 2, max_percdamp) + print(f" [sparsegpt] Cholesky failed, retrying with percdamp={damp_scale:.3f}") + continue + + losses = torch.zeros(self.rows, device=self.dev) + + for i1 in range(0, self.columns, blocksize): + i2 = min(i1 + blocksize, self.columns) + count = i2 - i1 + + W1 = W[:, i1:i2].clone() + Q1 = torch.zeros_like(W1) + Err1 = torch.zeros_like(W1) + Hinv1 = Hinv[i1:i2, i1:i2] + + if prune_n == 0: + # Unstructured: the OBS saliency w^2 / [H^-1]_ii^2, with the + # threshold chosen per block of columns. + saliency = W1.pow(2) / torch.diag(Hinv1).reshape(1, -1).pow(2) + thresh = torch.sort(saliency.flatten())[0][int(saliency.numel() * sparsity)] + mask1 = saliency <= thresh + else: + mask1 = torch.zeros_like(W1, dtype=torch.bool) + + for i in range(count): + w = W1[:, i] + d = Hinv1[i, i] + + if prune_n != 0 and i % prune_m == 0: + # Choose the n weights to drop inside this group of m. + group = (W1[:, i:i + prune_m].pow(2) / + torch.diag(Hinv1)[i:i + prune_m].reshape(1, -1).pow(2)) + mask1.scatter_(1, i + torch.topk(group, prune_n, dim=1, largest=False)[1], + True) + + q = w.clone() + q[mask1[:, i]] = 0 + Q1[:, i] = q + losses += (w - q).pow(2) / d.pow(2) + + # Propagate the error of column i onto the remaining + # columns of the block -- this is what makes SparseGPT + # stronger than magnitude pruning. + err1 = (w - q) / d + W1[:, i:] -= err1.unsqueeze(1) @ Hinv1[i, i:].unsqueeze(0) + Err1[:, i] = err1 + + W[:, i1:i2] = Q1 + W[:, i2:] -= Err1 @ Hinv[i1:i2, i2:] + + self.layer.weight.data.copy_(W.reshape_as(self.layer.weight).to(self.layer.weight.dtype)) + return (losses.sum() / 2).item() + + def free(self) -> None: + self.H = None + + +def hessian_bytes(layer: nn.Linear) -> int: + """Memory a :class:`SparseGPT` instance holds for ``layer`` (fp32 Hessian).""" + cols = layer.weight.data.shape[1] + return cols * cols * 4 diff --git a/compression/reasoning_aware_compression/rac/wanda.py b/compression/reasoning_aware_compression/rac/wanda.py new file mode 100644 index 000000000..93c9e0fe2 --- /dev/null +++ b/compression/reasoning_aware_compression/rac/wanda.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""Wanda pruning of a single linear layer. + +Adapted from `A Simple and Effective Pruning Approach for Large Language +Models `_ (Sun et al., ICLR 2024), +https://github.com/locuslab/wanda (MIT). + +Wanda drops the weights with the smallest :math:`|W_{ij}| \\cdot \\|X_j\\|_2`, +comparing weights *within each output row*. It needs only the per-input-channel +activation norm rather than a full Hessian, which makes it a cheap sanity check +that the RAC calibration set -- and not some SparseGPT-specific detail -- is +what drives the accuracy gains (Table 9 of the RAC paper). +""" + +from __future__ import annotations + +import torch +import torch.nn as nn + + +class Wanda: + """Accumulates :math:`\\|X_j\\|_2^2` per input channel, then prunes.""" + + def __init__(self, layer: nn.Linear): + self.layer = layer + self.dev = layer.weight.device + self.rows, self.columns = layer.weight.data.shape + self.scaler_row = torch.zeros(self.columns, device=self.dev, dtype=torch.float32) + self.nsamples = 0 + + @torch.no_grad() + def add_batch(self, inp: torch.Tensor) -> None: + if inp.dim() == 3: + inp = inp.reshape(-1, inp.shape[-1]) + elif inp.dim() == 1: + inp = inp.unsqueeze(0) + inp = inp.t().float() # (in_features, tokens) + + ntokens = inp.shape[1] + self.scaler_row *= self.nsamples / (self.nsamples + ntokens) + self.nsamples += ntokens + self.scaler_row += inp.pow(2).sum(dim=1) / self.nsamples + + @torch.no_grad() + def prune(self, sparsity: float, prune_n: int = 0, prune_m: int = 0, **_) -> float: + """Prune the layer in place and return the summed dropped saliency.""" + if self.nsamples == 0: + raise RuntimeError("Wanda.prune() called before any calibration batch.") + + W = self.layer.weight.data + metric = W.abs().float() * torch.sqrt(self.scaler_row).reshape(1, -1) + mask = torch.zeros_like(metric, dtype=torch.bool) + + if prune_n != 0: + for col in range(0, self.columns, prune_m): + group = metric[:, col:col + prune_m] + idx = torch.topk(group, prune_n, dim=1, largest=False)[1] + mask[:, col:col + prune_m].scatter_(1, idx, True) + else: + k = int(round(self.columns * sparsity)) + if k == 0: + return 0.0 + k = min(k, self.columns - 1) + idx = torch.topk(metric, k, dim=1, largest=False)[1] + mask.scatter_(1, idx, True) + + dropped = metric[mask].sum().item() + W[mask] = 0 + return dropped + + def free(self) -> None: + self.scaler_row = None diff --git a/compression/reasoning_aware_compression/requirements.txt b/compression/reasoning_aware_compression/requirements.txt new file mode 100644 index 000000000..0ae856125 --- /dev/null +++ b/compression/reasoning_aware_compression/requirements.txt @@ -0,0 +1,15 @@ +# Core: trace collection (HF backend) + pruning +torch>=2.1 +transformers>=4.45 +datasets>=2.19 +accelerate>=0.30 + +# Optional, pick the trace-collection backend you want: +# vllm>=0.6.0 # --backend vllm, fastest on a single node +# deepspeed>=0.14.0 # --backend deepspeed, ZeRO-Inference for 32B/70B models + +# Optional, evaluation of the pruned checkpoint (bash_script/eval_math500.sh): +# lighteval[math,vllm]>=0.8.0 + +# Optional, CPU-only unit tests: +# pytest>=7.0 diff --git a/compression/reasoning_aware_compression/tests/test_calibration.py b/compression/reasoning_aware_compression/tests/test_calibration.py new file mode 100644 index 000000000..a446aac23 --- /dev/null +++ b/compression/reasoning_aware_compression/tests/test_calibration.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""CPU-only tests for calibration-set construction. + +Run with ``pytest`` from ``compression/reasoning_aware_compression``. +""" + +import json +import sys +from pathlib import Path + +import pytest +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from rac.calibration import ( # noqa: E402 + build_calibration_samples, + chat_template_prompt, + load_trace_texts, + pack_token_windows, +) + + +class StubTokenizer: + """Whitespace tokenizer, so the tests never touch the network.""" + + eos_token_id = 0 + chat_template = None + + def __call__(self, text, add_special_tokens=False): + ids = [(abs(hash(word)) % 1000) + 1 for word in text.split()] + return type("Encoding", (), {"input_ids": ids})() + + def apply_chat_template(self, messages, tokenize=False, add_generation_prompt=True): + rendered = "".join(f"<|{m['role']}|>{m['content']}" for m in messages) + return rendered + "<|assistant|>" if add_generation_prompt else rendered + + +@pytest.fixture +def tokenizer(): + return StubTokenizer() + + +@pytest.fixture +def trace_file(tmp_path): + path = tmp_path / "traces.jsonl" + with path.open("w") as handle: + for i in range(20): + handle.write( + json.dumps({ + "prompt": " ".join(f"prompt{i}_{j}" for j in range(20)), + "completion": " ".join(f"cot{i}_{j}" for j in range(200)), + "prompt_tokens": 20, + "completion_tokens": 200, + }) + "\n") + return path + + +def test_pack_token_windows_shape(tokenizer): + texts = [" ".join(f"tok{i}_{j}" for j in range(100)) for i in range(50)] + samples = pack_token_windows(texts, tokenizer, nsamples=8, seqlen=64) + + assert len(samples) == 8 + for sample in samples: + assert sample.shape == (64, ) + assert sample.dtype == torch.long + + +def test_pack_token_windows_is_deterministic(tokenizer): + texts = [" ".join(f"tok{i}_{j}" for j in range(100)) for i in range(50)] + a = pack_token_windows(texts, tokenizer, nsamples=4, seqlen=32, seed=7) + b = pack_token_windows(texts, tokenizer, nsamples=4, seqlen=32, seed=7) + + assert all(torch.equal(x, y) for x, y in zip(a, b)) + + +def test_pack_token_windows_rejects_short_corpus(tokenizer): + with pytest.raises(ValueError, match="Calibration corpus yielded only"): + pack_token_windows(["one two three"], tokenizer, nsamples=4, seqlen=64) + + +def test_rac_source_includes_cot_tokens(trace_file, tokenizer): + """The whole method: RAC calibration must contain decode-time tokens.""" + prompt_only = list(load_trace_texts(str(trace_file), with_completion=False)) + with_cot = list(load_trace_texts(str(trace_file), with_completion=True)) + + assert len(prompt_only) == len(with_cot) == 20 + for prompt, full in zip(prompt_only, with_cot): + assert full.startswith(prompt) + assert len(full) > len(prompt) + assert "cot" in full and "cot" not in prompt + + +def test_build_calibration_samples_rac_vs_prompt(trace_file, tokenizer): + # 20 traces x (20 prompt + 200 CoT) tokens is plenty for a 1024-token budget. + rac = build_calibration_samples("rac", tokenizer, nsamples=16, seqlen=64, + traces=str(trace_file)) + assert len(rac) == 16 + + # The same traces hold 10x fewer prompt tokens, so the prompt-only baseline + # runs out at the same budget -- the CoT is where the tokens are. + with pytest.raises(ValueError, match="Calibration corpus yielded only"): + build_calibration_samples("prompt", tokenizer, nsamples=16, seqlen=64, + traces=str(trace_file)) + + +def test_build_calibration_samples_validates_inputs(tokenizer): + with pytest.raises(ValueError, match="requires --traces"): + build_calibration_samples("rac", tokenizer, nsamples=1, seqlen=8) + + with pytest.raises(ValueError, match="requires either --traces or --dataset"): + build_calibration_samples("prompt", tokenizer, nsamples=1, seqlen=8) + + with pytest.raises(ValueError, match="Unknown calibration source"): + build_calibration_samples("wikitext", tokenizer, nsamples=1, seqlen=8) + + +def test_chat_template_prompt(tokenizer): + assert chat_template_prompt(tokenizer, "2+2?", use_chat_template=False) == "2+2?" + # No chat_template on the stub tokenizer means the raw prompt is returned. + assert chat_template_prompt(tokenizer, "2+2?") == "2+2?" + + tokenizer.chat_template = "stub" + rendered = chat_template_prompt(tokenizer, "2+2?", system_prompt="be brief") + assert rendered == "<|system|>be brief<|user|>2+2?<|assistant|>" + + +def test_load_trace_texts_missing_file(tmp_path): + with pytest.raises(FileNotFoundError, match="Run collect_traces.py first"): + list(load_trace_texts(str(tmp_path / "nope.jsonl"))) diff --git a/compression/reasoning_aware_compression/tests/test_pruning.py b/compression/reasoning_aware_compression/tests/test_pruning.py new file mode 100644 index 000000000..70ea0c689 --- /dev/null +++ b/compression/reasoning_aware_compression/tests/test_pruning.py @@ -0,0 +1,192 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""CPU-only tests for the pruners and the block-by-block driver. + +A tiny randomly-initialised Qwen2 is built from a config, so the tests are +offline and run in a few seconds. Run with ``pytest`` from +``compression/reasoning_aware_compression``. +""" + +import sys +from pathlib import Path + +import pytest +import torch +from transformers import Qwen2Config, Qwen2ForCausalLM + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from rac import ( # noqa: E402 + magnitude_prune, + model_sparsity, + prunable_layers, + select_blocks_by_thirds, + sequential_prune, +) + +SEQLEN = 32 +NSAMPLES = 4 + + +@pytest.fixture +def model(): + torch.manual_seed(0) + config = Qwen2Config( + vocab_size=256, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=3, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=128, + torch_dtype="float32", + ) + tiny = Qwen2ForCausalLM(config) + tiny.eval() + return tiny + + +@pytest.fixture +def samples(): + generator = torch.Generator().manual_seed(1) + return [torch.randint(0, 256, (SEQLEN, ), generator=generator) for _ in range(NSAMPLES)] + + +def layer_sparsities(model, scope="all"): + return { + name: (layer.weight.data == 0).float().mean().item() + for name, layer in prunable_layers(model, scope=scope).items() + } + + +@pytest.mark.parametrize("method", ["sparsegpt", "wanda"]) +def test_sequential_prune_hits_target_sparsity(model, samples, method): + stats = sequential_prune(model, samples, method=method, sparsity=0.5, device="cpu") + + assert len(stats) == 3 * 7 # 3 blocks x (q,k,v,o,gate,up,down) + for name, sparsity in layer_sparsities(model).items(): + assert sparsity == pytest.approx(0.5, abs=0.05), name + + +def test_sequential_prune_semi_structured_2_of_4(model, samples): + sequential_prune(model, samples, method="sparsegpt", prune_n=2, prune_m=4, device="cpu") + + for name, layer in prunable_layers(model).items(): + W = layer.weight.data + groups = W.reshape(W.shape[0], -1, 4) + zeros_per_group = (groups == 0).sum(dim=2) + assert torch.all(zeros_per_group == 2), f"{name} is not 2:4 sparse" + + +def test_sequential_prune_respects_block_selection(model, samples): + # Prune only the first third of a 3-block model, i.e. block 0. + block_indices = select_blocks_by_thirds(3, [1]) + assert block_indices == [0] + + sequential_prune(model, samples, sparsity=0.5, block_indices=block_indices, device="cpu") + + sparsities = layer_sparsities(model) + assert all(v == pytest.approx(0.5, abs=0.05) for k, v in sparsities.items() + if k.startswith("model.layers.0.")) + assert all(v == 0.0 for k, v in sparsities.items() if not k.startswith("model.layers.0.")) + + +def test_sequential_prune_mlp_scope_leaves_attention_dense(model, samples): + sequential_prune(model, samples, sparsity=0.5, scope="mlp", device="cpu") + + for name, sparsity in layer_sparsities(model, scope="all").items(): + if "mlp" in name: + assert sparsity == pytest.approx(0.5, abs=0.05), name + else: + assert sparsity == 0.0, name + + +def test_sequential_prune_batched_matches_unbatched(model, samples): + import copy + + batched = copy.deepcopy(model) + sequential_prune(model, samples, sparsity=0.5, device="cpu", batch_size=1) + sequential_prune(batched, samples, sparsity=0.5, device="cpu", batch_size=NSAMPLES) + + for (name, a), (_, b) in zip(prunable_layers(model).items(), + prunable_layers(batched).items()): + assert torch.allclose(a.weight, b.weight, atol=1e-4), name + + +def test_sequential_prune_rejects_ragged_samples(model): + ragged = [torch.zeros(SEQLEN, dtype=torch.long), torch.zeros(SEQLEN + 1, dtype=torch.long)] + with pytest.raises(ValueError, match="same length"): + sequential_prune(model, ragged, sparsity=0.5, device="cpu") + + +def test_sequential_prune_rejects_unknown_method(model, samples): + with pytest.raises(ValueError, match="Unknown calibrated pruning method"): + sequential_prune(model, samples, method="alps", device="cpu") + + +def test_magnitude_prune(model): + magnitude_prune(model, sparsity=0.5) + assert model_sparsity(model) == pytest.approx(0.5, abs=0.02) + + +def test_pruned_model_still_runs(model, samples): + sequential_prune(model, samples, sparsity=0.5, device="cpu") + with torch.no_grad(): + logits = model(samples[0].reshape(1, -1)).logits + assert logits.shape == (1, SEQLEN, 256) + assert torch.isfinite(logits).all() + + +def test_calibration_data_changes_the_mask(model, samples): + """The premise of RAC: which tokens you calibrate on changes what is cut.""" + import copy + + other_generator = torch.Generator().manual_seed(99) + other_samples = [ + torch.randint(0, 256, (SEQLEN, ), generator=other_generator) for _ in range(NSAMPLES) + ] + + a, b = model, copy.deepcopy(model) + sequential_prune(a, samples, sparsity=0.5, device="cpu") + sequential_prune(b, other_samples, sparsity=0.5, device="cpu") + + masks_differ = any( + not torch.equal(la.weight == 0, lb.weight == 0) + for la, lb in zip(prunable_layers(a).values(), prunable_layers(b).values())) + assert masks_differ + + +def test_sparsegpt_handles_dead_input_channels(): + """A channel that never fires has a zero Hessian diagonal; it must not crash.""" + from rac.sparsegpt import SparseGPT + + torch.manual_seed(0) + layer = torch.nn.Linear(8, 4, bias=False) + activations = torch.randn(16, 8) + activations[:, 3] = 0.0 + + pruner = SparseGPT(layer) + pruner.add_batch(activations) + pruner.prune(0.5) + + assert torch.all(layer.weight.data[:, 3] == 0) + assert torch.isfinite(layer.weight.data).all() + + +@pytest.mark.parametrize("cls_name", ["SparseGPT", "Wanda"]) +def test_pruner_requires_calibration(cls_name): + import rac + + pruner = getattr(rac, cls_name)(torch.nn.Linear(8, 4, bias=False)) + with pytest.raises(RuntimeError, match="before any calibration batch"): + pruner.prune(0.5) + + +def test_select_blocks_by_thirds(): + assert select_blocks_by_thirds(9, [1]) == [0, 1, 2] + assert select_blocks_by_thirds(9, [1, 3]) == [0, 1, 2, 6, 7, 8] + assert select_blocks_by_thirds(9, [1, 2, 3]) == list(range(9)) + assert select_blocks_by_thirds(10, [2]) == [3, 4, 5] + + with pytest.raises(ValueError, match="Invalid third"): + select_blocks_by_thirds(9, [4]) From 95052f6ac42c29e45016bd9df03e6ab4142a70b7 Mon Sep 17 00:00:00 2001 From: Zhipeng Wang Date: Wed, 29 Jul 2026 06:48:50 -0700 Subject: [PATCH 2/2] Address review comments from delock and tohtana tohtana: include the upstream license files for the adapted pruners. Adds third_party_licenses/ with Wanda's MIT license (c) 2023 CMU Locus Lab and SparseGPT's Apache-2.0 text, and points the module docstrings and the README Attribution section at them. delock: - sequential.py: restore model.config.use_cache in a finally block so a failed pruning run does not leave the config mutated. - collect_traces.py: comment why n_completion counts non-pad tokens rather than the generated length. - collect_traces.py: all-reduce the running token total before the loop's exit check, so every ZeRO-Inference rank leaves on the same iteration instead of relying on the per-rank counts matching; add a cross-rank checksum guard that fails loudly if the ranks decoded different tokens. - modelutils.py: warn when find_linear_layers finds nothing (Conv1D-based architectures such as GPT-2/BLOOM), and fail in prune.py before calibration rather than saving a still-dense checkpoint. - README: "closely approximates" instead of "exactly", with a note that traces are re-tokenized when calibration windows are packed. Adds six regression tests (29 total, still CPU-only). --- .../reasoning_aware_compression/README.md | 21 +- .../collect_traces.py | 73 ++++++- .../reasoning_aware_compression/prune.py | 16 +- .../rac/modelutils.py | 16 ++ .../rac/sequential.py | 14 +- .../rac/sparsegpt.py | 3 +- .../reasoning_aware_compression/rac/wanda.py | 4 +- .../tests/test_collect_traces.py | 29 +++ .../tests/test_pruning.py | 45 ++++ .../third_party_licenses/LICENSE.sparsegpt | 201 ++++++++++++++++++ .../third_party_licenses/LICENSE.wanda | 21 ++ .../third_party_licenses/README.md | 12 ++ 12 files changed, 443 insertions(+), 12 deletions(-) create mode 100644 compression/reasoning_aware_compression/tests/test_collect_traces.py create mode 100644 compression/reasoning_aware_compression/third_party_licenses/LICENSE.sparsegpt create mode 100644 compression/reasoning_aware_compression/third_party_licenses/LICENSE.wanda create mode 100644 compression/reasoning_aware_compression/third_party_licenses/README.md diff --git a/compression/reasoning_aware_compression/README.md b/compression/reasoning_aware_compression/README.md index de5cb8b11..80b7e7124 100644 --- a/compression/reasoning_aware_compression/README.md +++ b/compression/reasoning_aware_compression/README.md @@ -54,9 +54,12 @@ recovers most of the loss: the same 7B model at 50% sparsity reaches 0.900 accur └──────────────────────────────────────┘ ``` -Teacher-forcing a sampled trace reproduces exactly the hidden states the model computed +Teacher-forcing a sampled trace closely approximates the hidden states the model computed while generating it, so phase II recovers the decode-time activations of Algorithm 1 in -the paper without re-running generation during pruning. +the paper without re-running generation during pruning. It is an approximation rather than +an identity because traces are stored as decoded text and re-tokenized when the +calibration windows are packed, so a few tokens near window boundaries can differ from +those originally sampled — negligible over a 1M-token calibration set. ## Layout @@ -77,6 +80,7 @@ compression/reasoning_aware_compression/ │ ├── collect_traces_zero_inference.sh # traces for 32B/70B via ZeRO-Inference │ └── eval_math500.sh # lighteval MATH-500 (accuracy + runtime) ├── tests/ # CPU-only unit tests (pytest) +├── third_party_licenses/ # upstream licenses for the adapted pruners └── requirements.txt ``` @@ -197,13 +201,18 @@ pytest tests ``` They cover calibration-window packing, the RAC-vs-prompt distinction, target sparsity for -SparseGPT/Wanda/magnitude, 2:4 patterns, block and scope selection, and batched-equals- -unbatched calibration on a tiny randomly-initialised Qwen2. +SparseGPT/Wanda/magnitude, 2:4 patterns, block and scope selection, batched-equals- +unbatched calibration on a tiny randomly-initialised Qwen2, and the failure paths +(unsupported architectures, config restored after a failed run). ## Attribution -- `rac/sparsegpt.py` adapts [IST-DASLab/sparsegpt](https://github.com/IST-DASLab/sparsegpt) (Apache-2.0). -- `rac/wanda.py` adapts [locuslab/wanda](https://github.com/locuslab/wanda) (MIT). +- `rac/sparsegpt.py` adapts [IST-DASLab/sparsegpt](https://github.com/IST-DASLab/sparsegpt) + (Apache-2.0); upstream license text in + [`third_party_licenses/LICENSE.sparsegpt`](third_party_licenses/LICENSE.sparsegpt). +- `rac/wanda.py` adapts [locuslab/wanda](https://github.com/locuslab/wanda) + (MIT, © 2023 CMU Locus Lab); upstream license text in + [`third_party_licenses/LICENSE.wanda`](third_party_licenses/LICENSE.wanda). - The RAC method and its reference implementation are by the paper's authors ([repo](https://github.com/RyanLucas3/Reasoning-Aware-Compression)); this example is a standalone reimplementation for DeepSpeedExamples. diff --git a/compression/reasoning_aware_compression/collect_traces.py b/compression/reasoning_aware_compression/collect_traces.py index 720a23ba1..4deb94ec0 100644 --- a/compression/reasoning_aware_compression/collect_traces.py +++ b/compression/reasoning_aware_compression/collect_traces.py @@ -239,6 +239,69 @@ def build_deepspeed_model(args): return engine.module +def _distributed_world() -> int: + """Number of ranks participating, or 1 when running unsharded.""" + try: + import torch.distributed as dist + except ImportError: + return 1 + if not (dist.is_available() and dist.is_initialized()): + return 1 + return dist.get_world_size() + + +def assert_ranks_agree(generated, device) -> None: + """Check that every rank decoded the same tokens. + + ZeRO-Inference shards parameters, not data: each rank runs the same batch + and must produce identical sequences, which is what makes it safe for rank 0 + alone to write the traces. Under a different sharding scheme (or a + rank-dependent seed) the ranks would diverge, and rank 0's file would + describe rollouts the other ranks never made -- so compare a cheap checksum + against rank 0 and fail loudly instead of writing inconsistent traces. + """ + import torch + import torch.distributed as dist + + if _distributed_world() < 2: + return + + checksum = torch.tensor( + [generated.shape[0], generated.shape[1], + int(generated.sum().item())], + dtype=torch.int64, + device=device, + ) + reference = checksum.clone() + dist.broadcast(reference, src=0) + if not torch.equal(checksum, reference): + raise RuntimeError( + f"Rank {dist.get_rank()} generated different tokens from rank 0 " + f"(shape/checksum {checksum.tolist()} vs {reference.tolist()}). Trace collection " + "assumes ZeRO-3 parameter sharding with replicated data, so every rank decodes the " + "same sequences; only rank 0 writes them.") + + +def _max_total_across_ranks(total: int, device) -> int: + """Largest running token count over all ranks. + + Every rank has to leave the generation loop on the same iteration: with + ``synced_gpus=True`` a rank that keeps going waits forever for peers that + stopped. The per-rank counts are expected to be equal -- see + :func:`assert_ranks_agree` -- but reducing them makes the exit collective + rather than a per-rank assumption. + """ + import torch + import torch.distributed as dist + + if _distributed_world() < 2: + return total + + counter = torch.tensor([total], dtype=torch.int64, device=device) + dist.all_reduce(counter, op=dist.ReduceOp.MAX) + return int(counter.item()) + + def generate_transformers(args, prompts: List[str], writer, use_deepspeed: bool) -> int: """Roll out with ``model.generate`` (HF or ZeRO-Inference backend).""" import torch @@ -281,18 +344,26 @@ def generate_transformers(args, prompts: List[str], writer, use_deepspeed: bool) synced_gpus=use_deepspeed, ) + if use_deepspeed: + assert_ranks_agree(generated, device) + prompt_len = batch["input_ids"].shape[1] completions = generated[:, prompt_len:] for i, sequence in enumerate(completions): prompt = chunk[i // args.num_generations] text = tokenizer.decode(sequence, skip_special_tokens=True) + # Number of *kept* tokens, not the generated length: generate() right-pads + # every sequence in the batch out to the longest one, and pad_token was + # aliased to eos_token above, so the terminal eos is not counted either. + # That matches the text actually written to the trace file, which is what + # the token budget is meant to track. n_completion = int((sequence != tokenizer.pad_token_id).sum()) n_prompt = int(batch["attention_mask"][i // args.num_generations].sum()) total += n_prompt + n_completion if is_writer_rank: writer(prompt, text, n_prompt, n_completion) - if total >= args.trace_tokens: + if _max_total_across_ranks(total, device) >= args.trace_tokens: break return total diff --git a/compression/reasoning_aware_compression/prune.py b/compression/reasoning_aware_compression/prune.py index 64afc622c..1816b68b0 100644 --- a/compression/reasoning_aware_compression/prune.py +++ b/compression/reasoning_aware_compression/prune.py @@ -35,6 +35,7 @@ layerwise_sparsity, magnitude_prune, model_sparsity, + prunable_layers, select_blocks_by_thirds, sequential_prune, ) @@ -129,8 +130,19 @@ def main() -> None: num_blocks = model.config.num_hidden_layers block_indices = select_blocks_by_thirds(num_blocks, thirds) target = (f"{args.prune_n}:{args.prune_m}" if args.prune_n else f"{args.sparsity:.0%}") - print(f"[rac] pruning {len(block_indices)}/{num_blocks} blocks to {target} " - f"with {args.pruning_method}, scope={args.scope}") + + # Fail here rather than after an hour of calibration: an architecture whose + # projections are not nn.Linear (GPT-2/BLOOM use Conv1D) would otherwise run + # the whole pass and save a checkpoint that is still dense. + targets = prunable_layers(model, scope=args.scope, block_indices=block_indices) + if not targets: + raise RuntimeError( + f"No prunable layers found in {args.model} with --scope {args.scope}. RAC prunes " + "nn.Linear projections inside decoder blocks; architectures that use " + "transformers Conv1D (GPT-2, BLOOM) are not supported.") + + print(f"[rac] pruning {len(targets)} layers in {len(block_indices)}/{num_blocks} blocks " + f"to {target} with {args.pruning_method}, scope={args.scope}") started = time.time() stats = [] diff --git a/compression/reasoning_aware_compression/rac/modelutils.py b/compression/reasoning_aware_compression/rac/modelutils.py index 8aab90277..1e8ff884e 100644 --- a/compression/reasoning_aware_compression/rac/modelutils.py +++ b/compression/reasoning_aware_compression/rac/modelutils.py @@ -4,6 +4,7 @@ from __future__ import annotations +import warnings from typing import Dict, Iterable, List, Optional, Tuple import torch @@ -75,6 +76,21 @@ def find_linear_layers(module: nn.Module, scope: str = "all") -> Dict[str, nn.Mo if scope == "mlp" and not any(k in name.lower() for k in MLP_KEYWORDS): continue found[name] = sub + + if not found: + # Silently returning nothing here would run the whole pruning pass and + # save a checkpoint that is still dense. The usual cause is an + # architecture whose projections are not ``nn.Linear`` -- GPT-2 and + # BLOOM use ``transformers.pytorch_utils.Conv1D`` -- or an MLP naming + # convention that ``--scope mlp`` does not recognise. + warnings.warn( + f"No prunable layers found in {type(module).__name__} with scope='{scope}'. " + f"RAC prunes {'/'.join(t.__name__ for t in PRUNABLE_TYPES)} sub-modules" + + (f" whose name contains one of {MLP_KEYWORDS}" if scope == "mlp" else "") + + "; this block will be left dense.", + RuntimeWarning, + stacklevel=2, + ) return found diff --git a/compression/reasoning_aware_compression/rac/sequential.py b/compression/reasoning_aware_compression/rac/sequential.py index 10abf3a52..23242fc4f 100644 --- a/compression/reasoning_aware_compression/rac/sequential.py +++ b/compression/reasoning_aware_compression/rac/sequential.py @@ -167,7 +167,20 @@ def sequential_prune( use_cache = getattr(model.config, "use_cache", False) model.config.use_cache = False + try: + stats = _prune_blocks(model, layers, selected, samples, pruner_cls, method, sparsity, + prune_n, prune_m, scope, device, blocksize, percdamp, batch_size) + finally: + # Restore even if pruning raises, so the caller is left with a usable + # model (and a usable config) to inspect. + model.config.use_cache = use_cache + return stats + +@torch.no_grad() +def _prune_blocks(model, layers, selected, samples, pruner_cls, method, sparsity, prune_n, + prune_m, scope, device, blocksize, percdamp, batch_size) -> List[dict]: + """The body of :func:`sequential_prune`; see it for the argument meanings.""" inps, block_kwargs = capture_block_inputs(model, samples, device) outs = torch.empty_like(inps) stats: List[dict] = [] @@ -220,7 +233,6 @@ def sequential_prune( _empty_cache(device) inps, outs = outs, inps - model.config.use_cache = use_cache return stats diff --git a/compression/reasoning_aware_compression/rac/sparsegpt.py b/compression/reasoning_aware_compression/rac/sparsegpt.py index 62e370ff8..a529fe0fa 100644 --- a/compression/reasoning_aware_compression/rac/sparsegpt.py +++ b/compression/reasoning_aware_compression/rac/sparsegpt.py @@ -5,7 +5,8 @@ Adapted from the reference implementation released with `SparseGPT: Massive Language Models Can Be Accurately Pruned in One-Shot `_ (Frantar & Alistarh, ICML 2023), -https://github.com/IST-DASLab/sparsegpt (Apache-2.0), and from the RAC +https://github.com/IST-DASLab/sparsegpt, Apache-2.0 licensed (upstream license +text reproduced in ``third_party_licenses/LICENSE.sparsegpt``), and from the RAC reference code at https://github.com/RyanLucas3/Reasoning-Aware-Compression. The algorithm itself is untouched by RAC: given the layer input activations diff --git a/compression/reasoning_aware_compression/rac/wanda.py b/compression/reasoning_aware_compression/rac/wanda.py index 93c9e0fe2..e7ebc130f 100644 --- a/compression/reasoning_aware_compression/rac/wanda.py +++ b/compression/reasoning_aware_compression/rac/wanda.py @@ -4,7 +4,9 @@ Adapted from `A Simple and Effective Pruning Approach for Large Language Models `_ (Sun et al., ICLR 2024), -https://github.com/locuslab/wanda (MIT). +https://github.com/locuslab/wanda, MIT licensed, Copyright (c) 2023 CMU Locus +Lab. The upstream license text is reproduced in +``third_party_licenses/LICENSE.wanda``. Wanda drops the weights with the smallest :math:`|W_{ij}| \\cdot \\|X_j\\|_2`, comparing weights *within each output row*. It needs only the per-input-channel diff --git a/compression/reasoning_aware_compression/tests/test_collect_traces.py b/compression/reasoning_aware_compression/tests/test_collect_traces.py new file mode 100644 index 000000000..97c2bac21 --- /dev/null +++ b/compression/reasoning_aware_compression/tests/test_collect_traces.py @@ -0,0 +1,29 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""CPU-only tests for the trace-collection helpers. + +Only the pieces that do not need a model are covered: the rank-coordination +helpers used by the ZeRO-Inference backend, which must degrade to no-ops when +the script runs in a single process. +""" + +import sys +from pathlib import Path + +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from collect_traces import _max_total_across_ranks, assert_ranks_agree # noqa: E402 + +CPU = torch.device("cpu") + + +def test_max_total_across_ranks_is_a_passthrough_without_distributed(): + assert _max_total_across_ranks(0, CPU) == 0 + assert _max_total_across_ranks(1_234_567, CPU) == 1_234_567 + + +def test_assert_ranks_agree_is_a_noop_without_distributed(): + generated = torch.randint(0, 100, (2, 8)) + assert assert_ranks_agree(generated, CPU) is None diff --git a/compression/reasoning_aware_compression/tests/test_pruning.py b/compression/reasoning_aware_compression/tests/test_pruning.py index 70ea0c689..bfae3ebba 100644 --- a/compression/reasoning_aware_compression/tests/test_pruning.py +++ b/compression/reasoning_aware_compression/tests/test_pruning.py @@ -17,6 +17,7 @@ sys.path.insert(0, str(Path(__file__).resolve().parents[1])) from rac import ( # noqa: E402 + get_decoder_layers, magnitude_prune, model_sparsity, prunable_layers, @@ -124,6 +125,50 @@ def test_sequential_prune_rejects_unknown_method(model, samples): sequential_prune(model, samples, method="alps", device="cpu") +@pytest.mark.parametrize("failing_kwargs, match", [ + ({"scope": "bogus"}, "Unknown scope"), + ({"method": "alps"}, "Unknown calibrated pruning method"), +]) +def test_sequential_prune_restores_use_cache_on_failure(model, samples, failing_kwargs, match): + """A failed pruning run must leave the config as it found it.""" + model.config.use_cache = True + + with pytest.raises(ValueError, match=match): + sequential_prune(model, samples, sparsity=0.5, device="cpu", **failing_kwargs) + + assert model.config.use_cache is True + + +def test_sequential_prune_restores_use_cache_on_success(model, samples): + model.config.use_cache = True + sequential_prune(model, samples, sparsity=0.5, device="cpu") + assert model.config.use_cache is True + + +def test_find_linear_layers_warns_when_nothing_is_prunable(model): + """Silently returning {} would save a checkpoint that is still dense.""" + import warnings + + from transformers.pytorch_utils import Conv1D + + from rac import find_linear_layers + + class Conv1DBlock(torch.nn.Module): + """Stand-in for a GPT-2/BLOOM block, whose projections are not nn.Linear.""" + + def __init__(self): + super().__init__() + self.attn = Conv1D(8, 8) + + with pytest.warns(RuntimeWarning, match="No prunable layers found"): + assert find_linear_layers(Conv1DBlock()) == {} + + # A supported block finds its layers and warns about nothing. + with warnings.catch_warnings(): + warnings.simplefilter("error") + assert len(find_linear_layers(get_decoder_layers(model)[0])) == 7 + + def test_magnitude_prune(model): magnitude_prune(model, sparsity=0.5) assert model_sparsity(model) == pytest.approx(0.5, abs=0.02) diff --git a/compression/reasoning_aware_compression/third_party_licenses/LICENSE.sparsegpt b/compression/reasoning_aware_compression/third_party_licenses/LICENSE.sparsegpt new file mode 100644 index 000000000..261eeb9e9 --- /dev/null +++ b/compression/reasoning_aware_compression/third_party_licenses/LICENSE.sparsegpt @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/compression/reasoning_aware_compression/third_party_licenses/LICENSE.wanda b/compression/reasoning_aware_compression/third_party_licenses/LICENSE.wanda new file mode 100644 index 000000000..606013019 --- /dev/null +++ b/compression/reasoning_aware_compression/third_party_licenses/LICENSE.wanda @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2023 CMU Locus Lab + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/compression/reasoning_aware_compression/third_party_licenses/README.md b/compression/reasoning_aware_compression/third_party_licenses/README.md new file mode 100644 index 000000000..40121c902 --- /dev/null +++ b/compression/reasoning_aware_compression/third_party_licenses/README.md @@ -0,0 +1,12 @@ +# Third-party licenses + +Two files in `rac/` are adapted from upstream pruning repositories. The original +license text of each is reproduced here verbatim, as their terms require. + +| File in this example | Adapted from | License | Text | +| --- | --- | --- | --- | +| `rac/sparsegpt.py` | [IST-DASLab/sparsegpt](https://github.com/IST-DASLab/sparsegpt) | Apache-2.0 | [`LICENSE.sparsegpt`](LICENSE.sparsegpt) | +| `rac/wanda.py` | [locuslab/wanda](https://github.com/locuslab/wanda) | MIT, © 2023 CMU Locus Lab | [`LICENSE.wanda`](LICENSE.wanda) | + +Everything else in this directory is original to DeepSpeedExamples and carries the +repository's Apache-2.0 license.