touchups based on remaining steps

Signed-off-by: Brian Dellabetta <[email protected]>

Author: brian-dellabetta

examples/quantization_w4a16/llama3_example.py

@@ -1,5 +1,4 @@
 import torch
-from compressed_tensors import force_cpu_offload
 from datasets import load_dataset
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
@@ -11,10 +10,9 @@
 
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
-    # device_map="auto",
+    device_map="cpu",
     torch_dtype="auto",
 )
-force_cpu_offload(model, execution_device=torch.device("cuda"))
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 
 # Select calibration dataset.
@@ -67,6 +65,7 @@ def tokenize(sample):
     recipe=recipe,
     max_seq_length=MAX_SEQUENCE_LENGTH,
     num_calibration_samples=NUM_CALIBRATION_SAMPLES,
+    oneshot_device=torch.device("cuda") if torch.cuda.is_available() else None,
 )
 
 # Confirm generations of the quantized model look sane.

src/llmcompressor/entrypoints/oneshot.py

@@ -2,6 +2,9 @@
 from datetime import datetime
 from typing import Optional
 
+import torch
+from compressed_tensors import force_cpu_offload
+from compressed_tensors.utils import get_execution_device
 from loguru import logger
 from torch.utils.data import DataLoader
 from transformers import PreTrainedModel
@@ -10,7 +13,11 @@
 from llmcompressor.core.session_functions import active_session
 from llmcompressor.datasets import get_calibration_dataloader
 from llmcompressor.entrypoints.utils import post_process, pre_process
-from llmcompressor.pipelines.registry import CalibrationPipeline
+from llmcompressor.pipelines import (
+    CalibrationPipeline,
+    LayerSequentialPipeline,
+    SequentialPipeline,
+)
 
 __all__ = ["Oneshot", "oneshot"]
 
@@ -186,6 +193,35 @@ def apply_recipe_modifiers(
         user_pipeline = self.dataset_args.pipeline
         modifiers = session.get_modifiers()
         pipeline = CalibrationPipeline.from_modifiers(modifiers, user=user_pipeline)
+
+        model_exec_device = get_execution_device(self.model)
+
+        # Sequential pipelines onload models layer by layer to minimize GPU memory usage
+        if isinstance(pipeline, (SequentialPipeline, LayerSequentialPipeline)):
+            # unless pure cpu run, throw warning if model lives on oneshot_device
+            if (
+                model_exec_device
+                == self.model_args.oneshot_device
+                != torch.device("cpu")
+            ):
+                logger.warning(
+                    f"Model device {model_exec_device} is the same as oneshot"
+                    " execution device. If you encounter OOM errors, consider"
+                    " loading the model up on CPU, so that more memory is available"
+                    " for the oneshot algorithm to run on GPU. Example available at"
+                    " examples/quantization_w4a16/llama3_example.py"
+                )
+
+            # set cpu offload for model
+            elif (
+                model_exec_device
+                == torch.device("cpu")
+                != self.model_args.oneshot_device
+            ):
+                force_cpu_offload(
+                    self.model, execution_devce=self.model_args.oneshot_device
+                )
+
         pipeline(self.model, calibration_dataloader, self.dataset_args)
 
         session.finalize()

src/llmcompressor/modifiers/awq/base.py

@@ -34,7 +34,6 @@
 __all__ = ["AWQModifier"]
 
 
-# TODO (Brian INFERENG-531) Add support for offloaded models
 class AWQModifier(Modifier, QuantizationMixin):
     """
     Implements the AWQ (Activation-Weighted Quantization) algorithm,

src/llmcompressor/pipelines/basic/pipeline.py

@@ -37,8 +37,6 @@ def __call__(
         :param dataloader: loads data for calibration
         :param dataset_args: dataset arguments relevant to pipelines
         """
-        # TODO: warn about cpu offloading
-
         model_device = get_execution_device(model)
 
         LifecycleCallbacks.calibration_epoch_start()

src/llmcompressor/pipelines/layer_sequential/pipeline.py

@@ -14,7 +14,10 @@
     to_next_layer_kwargs,
 )
 from llmcompressor.pipelines.registry import CalibrationPipeline
-from llmcompressor.pipelines.sequential.helpers import get_targets_from_modifiers
+from llmcompressor.pipelines.sequential.helpers import (
+    get_targets_from_modifiers,
+    keep_onload_context,
+)
 from llmcompressor.utils.helpers import DisableQuantization, calibration_forward_context
 
 if TYPE_CHECKING:
@@ -73,29 +76,34 @@ def __call__(
                 calib_desc = f"({layer_index + 1}/{num_layers}): Calibrating"
                 prop_desc = f"({layer_index + 1}/{num_layers}): Propagating"
 
-                # do a preliminary pass to trigger modifier hooks
-                for batch_idx in tqdm.tqdm(range(len(dataloader)), desc=calib_desc):
-                    inputs = intermediates.fetch(batch_idx)
-                    layer(**inputs)
-
-                LifecycleCallbacks.sequential_epoch_end()
-
-                # this pass does not trigger modifier hooks
-                # and is only used for capturing outputs from newly compressed modules
-                with HooksMixin.disable_hooks():
-                    for batch_idx in tqdm.tqdm(range(len(dataloader)), desc=prop_desc):
+                # reduce memory movement by keeping modules onloaded
+                with keep_onload_context():
+                    # do a preliminary pass to trigger modifier hooks
+                    for batch_idx in tqdm.tqdm(range(len(dataloader)), desc=calib_desc):
                         inputs = intermediates.fetch(batch_idx)
-                        output = layer(**inputs)
-
-                        if layer_index < num_layers - 1:
-                            next_layer = layers[layer_index + 1]
-                            output = to_next_layer_kwargs(output, next_layer)
-                            output = maybe_inject_pos_embeddings(
-                                output, next_layer, inputs
-                            )
-
-                            intermediates.delete(batch_idx)
-                            intermediates.update(batch_idx, output)
+                        layer(**inputs)
+
+                    LifecycleCallbacks.sequential_epoch_end()
+
+                    # this pass does not trigger modifier hooks
+                    # and is only used for capturing outputs from
+                    # newly compressed modules
+                    with HooksMixin.disable_hooks():
+                        for batch_idx in tqdm.tqdm(
+                            range(len(dataloader)), desc=prop_desc
+                        ):
+                            inputs = intermediates.fetch(batch_idx)
+                            output = layer(**inputs)
+
+                            if layer_index < num_layers - 1:
+                                next_layer = layers[layer_index + 1]
+                                output = to_next_layer_kwargs(output, next_layer)
+                                output = maybe_inject_pos_embeddings(
+                                    output, next_layer, inputs
+                                )
+
+                                intermediates.delete(batch_idx)
+                                intermediates.update(batch_idx, output)
 
             # redundant, finish any remaining compression
             LifecycleCallbacks.calibration_epoch_end()

src/llmcompressor/pipelines/sequential/pipeline.py

@@ -52,8 +52,6 @@ def __call__(
         """
         session = active_session()
 
-        # TODO: warn about not cpu offloading
-
         # prepare to trace subgraphs
         modifiers = session.get_modifiers()
         sequential_targets = get_targets_from_modifiers(modifiers, model)