[Performance] Batched calibration (#2054)

## Purpose ##
* Reduce calibration runtime by providing users with options to increase
performance
  * `batch_size` controls the batch size of calibration data
* `offload_sequential_activations` controls whether calibration data is
offloaded to the CPU between layers

## Prerequisites ##
* #2080
* #2081

## Changes ##

### Batched Calibration ###
* Add `batch_size` argument
* Change `data_collator` default from the default data collator to a
`"truncation"` collator
* The `data_collator_with_truncation` function truncates all samples to
the shortest length sample in the batch.
* Statistics about how many tokens are dropped using this method are in
the tables below
* The data collator can also be changed to "padding" instead to pad to
the longest length sample in the batch
* In order to reduce the amount of excess truncation/padding, default to
`LengthAwareSampler` which samples from the dataset such that samples
with similar batch lengths are batched together

Batch Size | Time | % Speedup | % Deleted
-- | -- | -- | --
Original (1) | 11m17 | N/A | 0.0
1 | 11m17 | 0.0 | 0.0
2 | 10m48 | 4.2 | 0.2
4 | 10m39 | 5.6 | 0.5
8 | 10m39 | 5.6 | 1.1
16 | 10m58 | 2.8 | 2.6
64 | 11m4 | 11.2 | 12.0
128 | 9m29 | 16.0 | 23.9
512 | 7m39 | 37.3 | 75.3
<!-- notionvc: 36bc5ab7-4968-4c6d-8f38-e5715769b9ba -->

* The speedup is relatively meager up until you start deleting
significant portions of the dataset via truncation

### Disable Offloading ###
* Add `offload_sequential_activations` argument, defaults to True (no
behavior change)
* Enabling this option increases throughput but also increases memory
usage

Batch Size | Time | % Speedup | % Deleted
-- | -- | -- | --
Original (1) | 11m17 | N/A | 0.0
1 | 10m14 | 9.3 | 0.0
2 | 9m46 | 13.4 | 0.2
4 | 9m36 | 14.9 | 0.5
8 | 9m48 | 13.1 | 1.1
16 | 9m26 | 16.3 | 2.6
32 | 9m27 | 16.2 | 5.8
128 | 8m34 | 24.0 | 23.9
512 | 6m40 | 40.9 | 75.3
<!-- notionvc: 3c954cd3-850c-412c-92b3-fa4cfa914be8 -->

* Memory requirement for 512 samples on Llama 8B is ~70Gb, which is
equivalent to batch size 128
* With this option enabled and batch size 32, calibration runtime is
less than 1s per layer (down from ~11s)
* This implies that the theoretical maximum speedup from reducing
calibration time alone is ~15% for this model + dataset

### Misc ###
* Fix examples
* Fixed examples where there's issues between model dtypes and processor
dtypes (Mixtral, Pixtral, Whisper)
* For multimodal models which use multimodal datasets, remove their data
collators, as the batch unwrapping is now done by
the`TextGenerationDataset`
* Remove `_mask_padding` from `IntermediatesCache`, as I do not believe
that this method is effective in masking padding tokens from hessian
calculations
* Fix AWQ
  * AWQ was hard coded to handle only batches of size 1

## Testing ##
### Evaluation Regression ###
Batch Size | Eval Score | Difference | % Deleted
-- | -- | -- | --
Original (1) | 0.6573 | 0.000 | 0.0
1 | 0.6513 | -0.6 | 0.0
2 | 0.6513 | -0.6 | 0.2
4 | 0.6657 | +0.8 | 0.5
8 | 0.6513 | -0.6 | 1.1
16 | 0.6672 | +1.0 | 2.6
64 | 0.6338 | -2.4 | 12.0
128 | 0.6603 | +0.3 | 23.9
512 | 0.6391 | -1.8 | 75.3
<!-- notionvc: c37244c6-0013-463e-9b1e-d82d6d78ebe1 -->

Deleting significant portions of the dataset (delete longer sequences
first) has a detrimental effect on recovery

### Modifiers ###
* GPTQ
  * Ran full regression tests, as shown above
* AWQ
  * Ran AWQ with batch size 32 and checked output sanity
* Quantization Modifier
  * Ran NVFP4 with batch size 10 and checked output sanity

### Calibration Regression Testing ###
I ran calibration for the following models (but did not evaluate
recovery)

The following model examples can calibrate without issue:
* Llama3
* Gemma3
* Internvl3
* Mllama
* Llama4

The following models had a bug where processor and model dtypes were
mismatched, but are now fixed by this PR:
* Mistral3
* Pixtral
* Whisper

The following models have an accelerate device offloading bug:
* Idefics3
* Phi3 Vision

The following model examples have an MoE replacement bug:
* qwen3-vl-30b-a3b-Instruct

## Future Work ##
While these options are a great place to start, the next step to improve
runtime is to allow multi-GPU compression, likely via torch.distributed
tensor parallelism

---------

Signed-off-by: Kyle Sayers <[email protected]>
Co-authored-by: Dipika Sikka <[email protected]>

Author: kylesayrs

examples/multimodal_audio/whisper_example.py

@@ -1,6 +1,10 @@
 import torch
 from datasets import load_dataset
-from transformers import WhisperForConditionalGeneration, WhisperProcessor
+from transformers import (
+    WhisperForConditionalGeneration,
+    WhisperProcessor,
+    default_data_collator,
+)
 
 from llmcompressor import oneshot
 from llmcompressor.modifiers.quantization import GPTQModifier
@@ -55,20 +59,27 @@ def process(sample):
         return_tensors="pt",
     )
 
-    inputs["input_features"] = inputs["input_features"].to(dtype=model.dtype)
+    # treat labels as calibration prefill
     inputs["decoder_input_ids"] = inputs["labels"]
     del inputs["labels"]
 
+    # strip extra dim added by multimodal processors
+    inputs = {key: value[0] for key, value in inputs.items()}
+
     return inputs
 
 
 ds = ds.map(process, remove_columns=ds.column_names)
 
 
-# Define a oneshot data collator for multimodal inputs.
-def data_collator(batch):
-    assert len(batch) == 1
-    return {key: torch.tensor(value) for key, value in batch[0].items()}
+# Patch: mismatch between processor and model dtype
+def data_collator(features):
+    for feature in features:
+        feature["input_features"] = torch.tensor(
+            feature["input_features"], dtype=model.dtype
+        )
+
+    return default_data_collator(features, return_tensors="pt")
 
 
 # Recipe

examples/multimodal_vision/gemma3_example.py

@@ -1,5 +1,4 @@
 import requests
-import torch
 from PIL import Image
 from transformers import AutoProcessor, Gemma3ForConditionalGeneration
 
@@ -13,17 +12,11 @@
 processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
 
 # Oneshot arguments
-DATASET_ID = "flickr30k"
-DATASET_SPLIT = {"calibration": "test[:512]"}
+BATCH_SIZE = 4
 NUM_CALIBRATION_SAMPLES = 512
 MAX_SEQUENCE_LENGTH = 2048
-
-
-# Define a oneshot data collator for multimodal inputs.
-def data_collator(batch):
-    assert len(batch) == 1
-    return {key: torch.tensor(value) for key, value in batch[0].items()}
-
+DATASET_ID = "flickr30k"
+DATASET_SPLIT = {"calibration": f"test[:{NUM_CALIBRATION_SAMPLES}]"}
 
 # Recipe
 recipe = [
@@ -41,14 +34,13 @@ def data_collator(batch):
 # Perform oneshot
 oneshot(
     model=model,
-    tokenizer=model_id,
+    processor=processor,
     dataset=DATASET_ID,
     splits=DATASET_SPLIT,
     recipe=recipe,
+    batch_size=BATCH_SIZE,
     max_seq_length=MAX_SEQUENCE_LENGTH,
     num_calibration_samples=NUM_CALIBRATION_SAMPLES,
-    trust_remote_code_model=True,
-    data_collator=data_collator,
 )
 
 # Confirm generations of the quantized model look sane.

examples/multimodal_vision/internvl3_example.py

@@ -37,20 +37,14 @@ def preprocess_and_tokenize(example):
         return_dict=True,
         return_tensors="pt",
     )
-    return inputs
-
 
-ds = ds.map(preprocess_and_tokenize)
+    # remove extra dim added by multimodal processors
+    inputs = {key: value[0] for key, value in inputs.items()}
 
+    return inputs
 
-def data_collator(batch):
-    assert len(batch) == 1
-    item = {key: value for key, value in batch[0].items()}
-    item["attention_mask"] = torch.tensor([item["attention_mask"]])
-    item["input_ids"] = torch.LongTensor([item["input_ids"]])
-
-    return item
 
+ds = ds.map(preprocess_and_tokenize, remove_columns=ds.column_names)
 
 # Recipe
 recipe = GPTQModifier(
@@ -68,7 +62,6 @@ def data_collator(batch):
     max_seq_length=MAX_SEQUENCE_LENGTH,
     num_calibration_samples=NUM_CALIBRATION_SAMPLES,
     trust_remote_code_model=True,
-    data_collator=data_collator,
 )
 
 # Save to disk compressed.

examples/multimodal_vision/llava_example.py

@@ -1,5 +1,4 @@
 import requests
-import torch
 from PIL import Image
 from transformers import AutoProcessor, LlavaForConditionalGeneration
 
@@ -19,12 +18,6 @@
 MAX_SEQUENCE_LENGTH = 2048
 
 
-# Define a oneshot data collator for multimodal inputs.
-def data_collator(batch):
-    assert len(batch) == 1
-    return {key: torch.tensor(value) for key, value in batch[0].items()}
-
-
 # Recipe
 recipe = [
     GPTQModifier(
@@ -44,7 +37,6 @@ def data_collator(batch):
     max_seq_length=MAX_SEQUENCE_LENGTH,
     num_calibration_samples=NUM_CALIBRATION_SAMPLES,
     trust_remote_code_model=True,
-    data_collator=data_collator,
     sequential_targets=["LlamaDecoderLayer"],
 )
 

examples/multimodal_vision/mistral3_example.py

@@ -4,7 +4,11 @@
 import requests
 import torch
 from PIL import Image
-from transformers import AutoProcessor, Mistral3ForConditionalGeneration
+from transformers import (
+    AutoProcessor,
+    Mistral3ForConditionalGeneration,
+    default_data_collator,
+)
 
 from llmcompressor import oneshot
 from llmcompressor.modifiers.quantization import GPTQModifier
@@ -27,17 +31,13 @@
 MAX_SEQUENCE_LENGTH = 2048
 
 
-# Define a oneshot data collator for multimodal inputs.
-def data_collator(batch):
-    assert len(batch) == 1
-    return {
-        key: (
-            torch.tensor(value)
-            if key != "pixel_values"
-            else torch.tensor(value, dtype=model.dtype)
+# Patch: mismatch between processor and model dtype
+def data_collator(features):
+    for feature in features:
+        feature["pixel_values"] = torch.tensor(
+            feature["pixel_values"], dtype=model.dtype
         )
-        for key, value in batch[0].items()
-    }
+    return default_data_collator(features, return_tensors="pt")
 
 
 # Recipe

examples/multimodal_vision/mllama_example.py

@@ -1,5 +1,4 @@
 import requests
-import torch
 from PIL import Image
 from transformers import AutoProcessor, MllamaForConditionalGeneration
 
@@ -19,12 +18,6 @@
 MAX_SEQUENCE_LENGTH = 2048
 
 
-# Define a oneshot data collator for multimodal inputs.
-def data_collator(batch):
-    assert len(batch) == 1
-    return {key: torch.tensor(value) for key, value in batch[0].items()}
-
-
 # Recipe
 recipe = [
     GPTQModifier(
@@ -44,7 +37,6 @@ def data_collator(batch):
     max_seq_length=MAX_SEQUENCE_LENGTH,
     num_calibration_samples=NUM_CALIBRATION_SAMPLES,
     trust_remote_code_model=True,
-    data_collator=data_collator,
     sequential_targets=["MllamaSelfAttentionDecoderLayer"],
 )
 

examples/multimodal_vision/pixtral_example.py

@@ -1,7 +1,11 @@
 import requests
 import torch
 from PIL import Image
-from transformers import AutoProcessor, LlavaForConditionalGeneration
+from transformers import (
+    AutoProcessor,
+    LlavaForConditionalGeneration,
+    default_data_collator,
+)
 
 from llmcompressor import oneshot
 from llmcompressor.modifiers.quantization import GPTQModifier
@@ -19,16 +23,13 @@
 MAX_SEQUENCE_LENGTH = 2048
 
 
-# Define a oneshot data collator for multimodal inputs.
-# NOTE: for transformers<4.48.0, please squeeze the first dimension of `pixel_values`
-# by appending `[0]` to the end of line 32
-def data_collator(batch):
-    assert len(batch) == 1
-    return {
-        "input_ids": torch.LongTensor(batch[0]["input_ids"]),
-        "attention_mask": torch.tensor(batch[0]["attention_mask"]),
-        "pixel_values": torch.tensor(batch[0]["pixel_values"]),
-    }
+# Patch: mismatch between processor and model dtype
+def data_collator(features):
+    for feature in features:
+        feature["pixel_values"] = torch.tensor(
+            feature["pixel_values"], dtype=model.dtype
+        )
+    return default_data_collator(features, return_tensors="pt")
 
 
 # Recipe
@@ -46,11 +47,11 @@ def data_collator(batch):
     tokenizer=model_id,
     dataset=DATASET_ID,
     splits={"calibration": f"{DATASET_SPLIT}[:{NUM_CALIBRATION_SAMPLES}]"},
+    data_collator=data_collator,
     recipe=recipe,
     max_seq_length=MAX_SEQUENCE_LENGTH,
     num_calibration_samples=NUM_CALIBRATION_SAMPLES,
     trust_remote_code_model=True,
-    data_collator=data_collator,
     sequential_targets=["MistralDecoderLayer"],
 )
 

src/llmcompressor/args/dataset_arguments.py

@@ -8,9 +8,7 @@
 """
 
 from dataclasses import dataclass, field
-from typing import Any, Callable
-
-from transformers import DefaultDataCollator
+from typing import Callable
 
 
 @dataclass
@@ -69,9 +67,27 @@ class CustomDatasetArguments(DVCDatasetArguments):
         },
     )
 
-    data_collator: Callable[[Any], Any] = field(
-        default_factory=lambda: DefaultDataCollator(),
-        metadata={"help": "The function to used to form a batch from the dataset"},
+    batch_size: int = field(
+        default=1,
+        metadata={
+            "help": (
+                "Calibration batch size. During calibration, LLM Compressor disables "
+                "lm_head output computations to reduce memory usage from large "
+                "batch sizes. Large batch sizes may result in excess padding or "
+                "truncation, depending on the data_collator"
+            )
+        },
+    )
+
+    data_collator: str | Callable = field(
+        default="truncation",
+        metadata={
+            "help": (
+                "The function to used to form a batch from the dataset. Can also "
+                "specify 'truncation' or 'padding' to truncate or pad non-uniform "
+                "sequence lengths in a batch. Defaults to 'padding'."
+            )
+        },
     )
 
 
@@ -126,8 +142,8 @@ class DatasetArguments(CustomDatasetArguments):
         default=512,
         metadata={"help": "Number of samples to use for one-shot calibration"},
     )
-    shuffle_calibration_samples: bool | None = field(
-        default=True,
+    shuffle_calibration_samples: bool = field(
+        default=False,
         metadata={
             "help": "whether to shuffle the dataset before selecting calibration data"
         },
@@ -142,7 +158,7 @@ class DatasetArguments(CustomDatasetArguments):
     )
     preprocessing_num_workers: int | None = field(
         default=None,
-        metadata={"help": "The number of processes to use for the preprocessing."},
+        metadata={"help": "The number of workers to use for dataset processing."},
     )
     pad_to_max_length: bool = field(
         default=True,
@@ -214,6 +230,14 @@ class DatasetArguments(CustomDatasetArguments):
             "definition"
         },
     )
+    offload_sequential_activations: bool = field(
+        default=True,
+        metadata={
+            "help": "Whether to offload intermediate activations between sequential "
+            "layers to the CPU. Disabling offloading is much faster, but uses "
+            "signficiantly more memory. Default is True."
+        },
+    )
     quantization_aware_calibration: bool = field(
         default=True,
         metadata={