add offload_sequential_activations

Signed-off-by: Kyle Sayers <[email protected]>

Author: kylesayrs

src/llmcompressor/args/dataset_arguments.py

@@ -229,6 +229,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={

src/llmcompressor/pipelines/sequential/pipeline.py

@@ -66,7 +66,6 @@ def __call__(
         # prepare to trace subgraphs
         modifiers = session.lifecycle.recipe.modifiers
         sequential_targets = get_sequential_targets(modifiers, model, dataset_args)
-
         ignore = dataset_args.tracing_ignore
 
         # trace subgraphs
@@ -90,7 +89,11 @@ def __call__(
                 stack.enter_context(DisableQuantization(model))
 
             # prepare intermediates cache
-            activations = IntermediatesCache.from_dataloader(dataloader, model_device)
+            cache_offload = dataset_args.offload_sequential_activations
+            offload_device = torch.device("cpu") if cache_offload else None
+            activations = IntermediatesCache.from_dataloader(
+                dataloader, model_device, offload_device=offload_device
+            )
 
             for subgraph_index, subgraph in enumerate(subgraphs):
                 # prepare tqdm description texts