MaskLLM
NVIDIA Research · Coding · Other
MaskLLM is an open-source pruning method from NVIDIA Research that learns N:M sparsity masks to shrink large language models, cutting inference cost while keeping accuracy close to the dense model.
About MaskLLM
What Is MaskLLM
MaskLLM is a research method for compressing large language models. It belongs to the family of model pruning tools that remove parameters the model does not need, and it does this through semi-structured sparsity, also called N:M sparsity. The idea is simple: in every block of M consecutive weights, keep only N of them. A 2:4 pattern keeps two of every four weights, which halves the size that hardware has to move.
What sets MaskLLM apart is how it picks which weights to keep. Older methods rely on a fixed rule, like throwing out the smallest weights or using a statistical estimate. MaskLLM instead treats the choice of mask as something the model learns. It models the N:M patterns as a probability distribution and trains it with a sampling trick called Gumbel Softmax, so the model can find better masks on a large dataset. The weights themselves stay frozen while the masks are learned, which is what makes the 2:4 sparsity "lossless" for a given task.
The practical draw sits in the numbers. According to NVIDIA, the dense LLaMA-2 7B model scores a perplexity of 5.12 on Wikitext, while earlier pruning methods push it past 10. MaskLLM lands at 6.72 to 6.78 by only learning the masks, with the weights frozen. The project is open-source under an NVIDIA license, so researchers and engineers can run it themselves. Running it well still takes real GPU hardware and machine-learning skills.
Getting Started
- Install the Python packages, then run the provided script to download a pre-computed mask from Hugging Face
- Load your target model, such as LLaMA-2 7B or LLaMA-3 8B, in PyTorch
- Apply the downloaded mask and evaluate accuracy with the included perplexity script
- For custom work, set up the Megatron-LM environment, prepare data and train a new mask on your own domain
Product Information
A quick look at MaskLLM's pricing, supported platforms, and performance.
Best for
The users, tasks, and scenarios where this tool fits best.
Users
- Machine learning researchers
- ML engineers deploying LLMs
- Graduate students and academics
- Framework and tooling developers
Tasks
- Pruning large language models
- One-shot compression
- Domain-specific sparsity
- Sparsity transfer
- Benchmarking pruning methods
Scenarios
- A team serving a 7B model on limited GPUs
- Research that must reproduce published numbers
- A model that works well in general but slips in a niche domain
- Building a compression pipeline for a model family
- Studying privacy or licensing limits
Key features
Learnable N:M Masks
Instead of scoring weights by hand, MaskLLM learns the mask itself. It treats the set of valid N:M patterns as a distribution and samples from it during training, so the model can converge on the best masks for a dataset. That learnable step is the core difference from one-shot methods, which pick a mask once and stop.
Gumbel Softmax Sampling
The method borrows a technique from discrete sampling to make the mask choice trainable end to end. Gumbel Softmax lets gradients flow through what would otherwise be a hard, all-or-nothing pick, which is what enables the distribution over N:M masks to be optimized on real data rather than estimated.
Pre-computed Masks for Hugging Face Models
The project ships ready-made masks so you can try it without training anything. There are published masks for LLaMA-2 7B, LLaMA-3 8B and LLaMA-3.1 8B, each trained on a C4 subset, and a script downloads them from Hugging Face and applies them to the official models for evaluation.
Semi-Structured 2:4 Sparsity
MaskLLM targets N:M sparsity, most often the 2:4 pattern that keeps two non-zero values in every four consecutive weights. This format matters because it maps onto hardware acceleration, so the smaller model does not just take less space, it can also run faster on the right GPU.
Sparse Pattern Transfer
The same probability model that learns a mask can move it across domains. Once a mask distribution is trained, it can be adapted to a new task or dataset with less work than training from zero, which the authors call transferability. That makes it useful when a model has to stay sharp in more than one setting.
Megatron-LM Training Pipeline
For full training runs, MaskLLM includes a Megatron-LM workflow that spreads a LLaMA model across multiple GPUs. The repository provides scripts to convert checkpoints, tokenize data, generate a prior mask from SparseGPT, Wanda or magnitude pruning, train the mask, and trim the checkpoint for release.
Pros and cons
Pros
- Learns masks from data instead of relying on a fixed importance rule, which improves accuracy at the same sparsity
- Publishes code, weights and pre-computed masks under an open license, so results can be reproduced
- Supports transfer learning, so one mask distribution can be reused across tasks and domains
- Works with popular models like LLaMA-2, LLaMA-3 and LLaMA-3.1 through Hugging Face and Megatron-LM
- Keeps weights frozen, so the learned 2:4 sparsity is lossless for the task it was tuned on
Cons
- Training a new mask takes a multi-GPU setup, and the repository points to roughly 40GB of memory per GPU for full runs
- It is a research artifact, not a product, so there is no support desk, no hosted service and no API
- Out-of-the-box masks ship for a few LLaMA models, so other architectures need you to train your own
- The benefit shows up on hardware that supports 2:4 acceleration, which limits what older or non-NVIDIA GPUs get from it
Frequently asked questions
MaskLLM is a learnable pruning method that adds N:M sparsity to large language models. Instead of scoring each weight, it learns which weights to keep by modeling the masks as a probability distribution. NVIDIA Research published it at NeurIPS 2024, and the MaskLLM code is open-source on GitHub.
Related content
Explore related tools, skills, and articles for MaskLLM.
MaskLLM Alternatives
Udon
Udon · Coding · ProductivityUdon turns an Apple Silicon Mac into a home server you run from a browser, with container controls, file shares, and an optional AI sysadmin.
offstage
offstage cloud · Coding · Productivityoffstage gives AI coding agents and CI jobs a real Mac: a fresh macOS account on Apple silicon, driven over an API and billed by the second.
BotBus
Chengdu Tuyang Tusen Technology Co., Ltd. · Coding · ProductivityBotBus lets you control AI coding agents from your phone and watch: approve commands, start tasks and read results while the code stays on your desktop.