SparseForge

Official code for "SparseForge: Efficient LLM Semi-Structured Pruning via Hessian-Guided Soft-Mask Tempering" Paper: arXiv  

SparseForge is a post-training framework for semi-structured LLM sparsification that optimizes the sparsity mask directly — rather than scaling up retraining tokens — to achieve efficient sparse recovery. It combines Hessian-aware importance estimation with progressive tempering of soft masks into hardware-executable 2:4 patterns.

Key result: On LLaMA-2-7B under 2:4 sparsity, SparseForge reaches 57.27% average zero-shot accuracy with only 5B retraining tokens — surpassing the dense model (56.43%) and approaching a 40B-token prior SOTA (57.52%) with ~8× fewer tokens.

Post-submission experiment pack. Unified baseline reevaluations, fixed-support recovery curves, reviewer-facing reporting notes, and a new-cluster continuation checklist are available in experiments/rebuttal_2026/.

Highlights

• Soft-mask optimization — treats the sparsity mask as a continuous, optimizable variable; avoids premature hard decisions that plague one-shot pruning.
• Hessian-guided scoring — uses stochastic Hutchinson estimation to capture deletion sensitivity under grouped (2:4) competition.
• Progressive tempering — gradually shapes soft masks into deployable binary patterns, closing the soft-to-hard gap without abrupt accuracy drops.
• SLoRB (Sparse Low-Rank Bypass) — lightweight correction modules for residual capacity.
• Multi-architecture — supports LLaMA / OPT / GPT-2 / Qwen / Mistral / DeepSeek-MoE / Hunyuan.
• Block-sparse-16 Triton kernels in addition to the classic 2:4 N:M pattern.
• Optional FSDP + DeepSpeed sharding for 7B+ scale training.

---

Repository Layout

SparseForge/
├── sparseforge/             # Modular Python package (clean API)
│   ├── __init__.py
│   ├── cli.py              # Argparse builders
│   ├── data_pipeline.py    # AsyncDataPrefetcher / get_batch
│   ├── distributed.py      # Multi-node barrier / debug logging
│   ├── optim_utils.py      # Optimizer construction, LR schedule
│   ├── eval_utils.py       # estimate_loss
│   ├── state.py            # TrainState dataclass
│   ├── checkpoint.py       # Resume / save logic
│   ├── training_loop.py    # Main training loop interface
│   └── model_builders/     # Per-architecture model builders
│       ├── llama.py
│       └── universal.py
│
├── legacy/                  # Monolithic single-file entry points (see below)
│   ├── main_llama.py       # LLaMA-2-7B training (all-in-one)
│   ├── main_universal.py   # Multi-architecture training (all-in-one)
│   ├── eval_wiki_ppl.py    # Standalone evaluation script
│   ├── sparse_modeling.py  # SparseLinear / Distill_Model
│   ├── utils.py            # Mask ops, Hessian, penalties, calibration
│   ├── model_factory.py    # Auto-dispatch model creation
│   ├── model_*.py          # Per-architecture model adapters
│   ├── adamw.py            # Mask-aware AdamW optimizer
│   └── triton_block_sparse.py  # Triton block-sparse kernels
│
├── channel_pruning/         # Structured channel-pruning utilities
├── data/                    # Dataset download & preparation scripts
├── configs/                 # DeepSpeed config + hostfile template
├── scripts/                 # Cluster launcher + data preparation helpers
├── assets/                  # Figures
│
├── train_llama.sh           # Launcher → legacy/main_llama.py
├── train_universal.sh       # Launcher → legacy/main_universal.py
├── eval_wiki_ppl.sh         # Launcher → legacy/eval_wiki_ppl.py
├── requirements.txt
└── README.md

---

Installation

git clone https://github.com/tsinghua-ideal/SparseForge.git
cd SparseForge

# 1. PyTorch (adjust the CUDA version to your system)
pip install torch --index-url https://download.pytorch.org/whl/cu121

# 2. Remaining dependencies
pip install -r requirements.txt

Tested with PyTorch 2.1+, CUDA 12.1, and 8×H800 / 8×A100 GPUs.

---

Alternative: Legacy Monolithic Scripts

The legacy/ directory contains the original all-in-one training scripts that were used to produce all paper results. Each script embeds the full training loop, model construction, mask scheduling, and evaluation inline — no external package dependencies beyond PyTorch and HuggingFace.

This is the battle-tested, numerically verified code path. If you just want to reproduce the paper numbers without touching the modular sparseforge/ package, use the shell launchers directly — they already point to legacy/:

bash train_llama.sh 1 0          # single-node LLaMA training
bash train_universal.sh 1 0      # single-node universal training
bash eval_wiki_ppl.sh 1 0        # single-node evaluation

See legacy/README.md for details.

---

Data Preparation

C4 (training corpus)

bash data/download_c4.sh                 # raw shards
# Per-tokenizer pre-tokenisation (one of the following, depending on the model):
python data/prepare_instruct.py          # instruction-style preprocessing
# or see scripts/prepare_mixed_c4_based.py for per-tokenizer binarization

Each tokenizer produces an isolated binary directory (e.g. data/c4_llama/, data/c4_qwen/, ...) which is referenced by --dataset c4_${MODEL_TYPE}.

WikiText-2 (evaluation)

python data/download_wikitext.py

Pre-trained models

python data/download_hf_model.py --repo NousResearch/Llama-2-7b-hf --out models/Llama--Llama2-7b

---

Training

Launch Scripts Overview

Controller mode (recommended): bash train_llama.sh <NNODES> <IDX1> ... <IDXN> — selects nodes from the pool defined in cluster_launcher.sh and auto-launches via SSH.

Legacy mode: bash train_llama.sh <MASTER_IP> <NODE_RANK> <NNODES> — run manually on each node.

LLaMA-2-7B (2:4 N:M sparsity, Hutchinson Hessian)

bash train_llama.sh

Single-node (default) uses deepspeed --num_gpus 8. For multi-node:

# Edit configs/hosts.txt with your <MASTER_IP> / <WORKER_IP>.
NNODES=2 NODE_RANK=0 MASTER_ADDR=<MASTER_IP> bash train_llama.sh

Use USE_FSDP_FULLY_SHARDED=1 bash train_llama.sh to train with PyTorch FSDP instead.

Universal trainer (OPT / Qwen / Mistral / DeepSeek-MoE / Hunyuan / GPT-2)

Open train_universal.sh and uncomment the desired model block, e.g.

STUDENT_MODEL="models/Qwen--Qwen3-1.7b"
TEACHER_MODEL="models/Qwen--Qwen3-1.7b"
MODEL_TYPE="qwen"
MASK_TYPE="nm_2_4"            # or "block_sparse16" / "unstructured"

Then:

bash train_universal.sh

Key arguments (shared by both entry points)

See legacy/main_universal.py --help for the full list.

---

Evaluation

# Evaluate a trained checkpoint on WikiText-2 and zero-shot benchmarks.
CKPT_PATH=outputs/.../model.pt \
MODEL_PATH=models/Qwen--Qwen3-1.7b \
bash eval_wiki_ppl.sh

# Inspect checkpoint sparsity
python legacy/check_sparsity.py --ckpt outputs/.../model.pt

eval_wiki_ppl.sh will optionally run lm_eval on boolq, rte, hellaswag, winogrande, arc_easy, arc_challenge, openbookqa with RUN_LM_EVAL=true.

---

Reproducing the Paper

Default hyper-parameters in the provided .sh scripts match the paper setup.

---

Results at a Glance

Cross-Model Summary (2:4 sparsity, mean zero-shot accuracy %)

LLaMA-2-7B Comparison (2:4 sparsity)

† The original table mixed checkpoint provenance for perplexity. The archived 5B checkpoint has a same-checkpoint WikiText-2 PPL of approximately 6.22 under the current evaluation setup. See the experiment pack for the full audit. | CAST† (40B tokens) | 40B | 57.52% | 5.21 |

---

Citation

If you find SparseForge useful, please cite:

@misc{hanzuo2026sparseforgeefficientsemistructuredllm,
      title={SparseForge: Efficient Semi-Structured LLM Sparsification via Annealing of Hessian-Guided Soft-Mask}, 
      author={Liu Hanzuo and Chaofan Lin and Weixuan Sun and Yulong Wang and Key and Rayying and Mingyu Gao},
      year={2026},
      eprint={2605.06402},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2605.06402}, 
}

---

Acknowledgements

SparseForge builds upon nanoGPT, lm-evaluation-harness, HuggingFace transformers / datasets, PyTorch FSDP, DeepSpeed and Triton. We thank the authors and maintainers of these projects.

License

Released under the Apache License 2.0.