[Performance] Sequential onloading (#1263)

# Sequential Onloading #
<p align="center"><img width="403" alt="Screenshot 2025-06-05 at 22 53
01"
src="https://github.com/user-attachments/assets/ffd610ac-c511-4dc1-b858-b0ed2bf95193"
/></p>

```
(25/33): Calibrating:   0%|                                                                                                                                  | 0/512 [00:00<?, ?it/s]
<class 'transformers.models.llama.modeling_llama.LlamaRMSNorm'>.weight -> cuda
<class 'torch.nn.modules.linear.Linear'>.weight -> cuda
<class 'torch.nn.modules.linear.Linear'>.weight_scale -> cuda
<class 'torch.nn.modules.linear.Linear'>.weight_zero_point -> cuda
...
(25/33): Calibrating: 100%|█████| 512/512 [00:23<00:00, 21.91it/s]
2025-06-03T17:29:15.536963-0400 | compress_modules | INFO - Quantizing model.layers.24.self_attn.q_proj using 512 samples
2025-06-03T17:29:17.328720-0400 | compress | METRIC - time 1.79s
2025-06-03T17:29:17.329265-0400 | compress | METRIC - error 8948.54
2025-06-03T17:29:17.329781-0400 | compress | METRIC - GPU 0 | usage: 5.41% | total memory: 85 GB
2025-06-03T17:29:17.330248-0400 | compress | METRIC - Compressed module size: 33.947648 MB
...
(25/33): Propagating: 100%|█████| 512/512 [00:03<00:00, 131.16it/s]
<class 'transformers.models.llama.modeling_llama.LlamaRMSNorm'>.weight -> meta
<class 'torch.nn.modules.linear.Linear'>.weight -> meta
<class 'torch.nn.modules.linear.Linear'>.weight_scale -> meta
<class 'torch.nn.modules.linear.Linear'>.weight_zero_point -> meta
...
```

## Purpose ##
* Reduce hardware requirements for calibrating large models
* Reduce runtime caused by excess device movement when calibrating
offloaded models

## Prerequisites ##
* neuralmagic/compressed-tensors#354
* neuralmagic/compressed-tensors#355
* neuralmagic/compressed-tensors#356
* neuralmagic/compressed-tensors#357

## Related Issues ##
* Resolves #1383
* Resolves #1228
* Resolves #1122
* Resolves #1078
* Resolves #1216
* Resolves #1483

## Changes ##
* Keep layer parameters onloaded during the entire sequential
calibration + compression + propagation step
* This is achieved through the `keep_onload_context`, which disables
offloaded until the context is exited
* Dispatch model within each calibration pipeline
* Sequential pipeline offloads the model to CPU, and executes on the
first cuda device
* Deprecate passing sequential_targets via modifiers, instead prefer
passing via oneshot argument
* Use sequential pipeline as default pipeline (basic pipeline is never
used)
* Deprecate passing sequential_targets via modifiers, instead prefer
passing via oneshot argument
* Dispatch model before sample generation
* The model is dispatched exactly as it would be if it was loaded with
`device_map="auto"`

### Examples ###
* Models are loaded onto cpu before oneshot (rather than being
dispatched across GPUs)
* Model is reloaded from disk in order to redispatch onto "auto" device
map
* In my opinion, this is a better flow anyways, since models can raise
errors / take a very long time during generation, which can cause the
entirely compression job to go to waste
* The alternative is to either call `accelerate.remove_hooks(model)` and
`accelerate.dispatch_model(model)` before generating, or get rid of
sample generation entirely. One of these may be required if
compressed_linear isn't reliable enough to add to our examples

<details><summary>New example script</summary>

```python3
from transformers import AutoModelForCausalLM

from llmcompressor.modifiers.quantization import GPTQModifier
from llmcompressor.transformers import oneshot
from llmcompressor.utils.dev import dispatch_for_generation

# Load model (on cpu)
model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto")  # model is loaded on cpu
tokenizer = AutoTokenizer.from_pretrained(model_id)

# Define recipe
recipe = GPTQModifier(targets="Linear", scheme="W4A16", ignore=["lm_head"])

# Apply oneshot (model execution device is set to cuda, model stays on cpu)
oneshot(
    model=model,
    dataset="ultrachat_200k",
    recipe=recipe,
    max_seq_length=2048,
    num_calibration_samples=512,
)

# Perform sample generation
print("\n\n")
print("========== SAMPLE GENERATION ==============")
dispatch_for_generation(model)
input_ids = tokenizer("Hello my name is", return_tensors="pt").input_ids.to("cuda")
output = model.generate(input_ids, max_new_tokens=100)
print(tokenizer.decode(output[0]))
print("==========================================\n\n")

# Save to disk before generating
SAVE_DIR = model_id.split("/")[1] + "-W4A16-G128"
model.save_pretrained(SAVE_DIR, save_compressed=True)
tokenizer.save_pretrained(SAVE_DIR)
```
</details>

## Testing ##
* Calibrated and GPTQ-compressed one layer of Deepseek-V3 with a single
H100 in 50 seconds
  * 4.5x Improvement over original 236 seconds
* Peak memory of ~40 GB, which can be further reduced by increasing the
granularity of sequential targets
* Not offloading activations did not result in a performance improvement
* TODO: Test all example models can be reloaded and run

---------

Signed-off-by: Kyle Sayers <[email protected]>
Signed-off-by: Brian Dellabetta <[email protected]>
Co-authored-by: Brian Dellabetta <[email protected]>

Author: kylesayrs

examples/awq/README.md

@@ -18,11 +18,7 @@ recipe = [
 To use your own model, start with an existing example change the `model_id` to match your own model stub.
 ```python
 model_id = "path/to/your/model"
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    device_map="auto",
-    torch_dtype="auto",
-)
+model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto")
 ```
 
 ## Adding Mappings ##

examples/awq/llama_example.py

@@ -7,9 +7,7 @@
 # Select model and load it.
 MODEL_ID = "meta-llama/Meta-Llama-3-8B-Instruct"
 
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID, device_map="auto", torch_dtype="auto"
-)
+model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype="auto")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
 
 # Select calibration dataset.

examples/awq/qwen3_moe_example.py

@@ -3,13 +3,12 @@
 
 from llmcompressor import oneshot
 from llmcompressor.modifiers.awq import AWQModifier
+from llmcompressor.utils import dispatch_for_generation
 
 # Select model and load it.
 MODEL_ID = "Qwen/Qwen3-30B-A3B"
 
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID, device_map="auto", torch_dtype="auto"
-)
+model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype="auto")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
 
 # Select calibration dataset.
@@ -71,6 +70,7 @@ def tokenize(sample):
 # Confirm generations of the quantized model look sane.
 print("\n\n")
 print("========== SAMPLE GENERATION ==============")
+dispatch_for_generation(model)
 input_ids = tokenizer("Hello my name is", return_tensors="pt").input_ids.to("cuda")
 output = model.generate(input_ids, max_new_tokens=100)
 print(tokenizer.decode(output[0]))

examples/big_models_with_accelerate/README.md

@@ -0,0 +1,12 @@
+## Big Modeling with Sequential Onloading ##
+### What is Sequential Onloading? ###
+Sequential onloading is a memory-efficient approach for compressing large language models (LLMs) using only a single GPU. Instead of loading the entire model into memory—which can easily require hundreds of gigabytes—this method loads and compresses one layer at a time. The outputs are offloaded before the next layer is processed, dramatically reducing peak memory usage while maintaining high compression fidelity.
+
+<p align="center">
+    <img src="assets/sequential_onloading.png"/>
+</p>
+
+For more information, see the [RedHat AI blog post](https://developers.redhat.com/articles/2025/05/09/llm-compressor-optimize-llms-low-latency-deployments#generalizing_to_multimodal_and_moe_architectures) or the [LLM Compressor Office Hours Recording](https://www.youtube.com/watch?v=GrhuqQDmBk8).
+
+### Using Sequential Onloading ###
+Sequential onloading is enabled by default within LLM Compressor. To disable sequential onloading, add the `pipeline="basic"` argument to the LLM Compressor `oneshot` function call.

examples/big_models_with_accelerate/cpu_offloading_fp8.py

@@ -21,11 +21,7 @@ This directory contains example scripts for quantizing a variety of audio langua
 To use your own multimodal modal, start with an existing example change the `model_id` to match your own model stub.
 ```python3
 model_id = "path/to/your/model"
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    device_map="auto",
-    torch_dtype="auto",
-)
+model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto")
 ```
 
 ## Customizing GPTQModifier Parameters ##