implement oneshot_device, pipeline warnings

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

Author: kylesayrs

src/llmcompressor/args/model_arguments.py

@@ -81,7 +81,7 @@ class ModelArguments:
         metadata={"help": "Whether to compress sparse models during save"},
     )
     oneshot_device: Optional[str] = field(
-        default="cuda:0",
+        default="cuda",
         metadata={"help": "Device to run oneshot calibration on"},
     )
     model_revision: str = field(

src/llmcompressor/entrypoints/oneshot.py

@@ -2,9 +2,6 @@
 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
@@ -13,11 +10,7 @@
 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 import (
-    CalibrationPipeline,
-    LayerSequentialPipeline,
-    SequentialPipeline,
-)
+from llmcompressor.pipelines import CalibrationPipeline
 
 __all__ = ["Oneshot", "oneshot"]
 
@@ -193,35 +186,6 @@ 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/entrypoints/utils.py

@@ -3,6 +3,8 @@
 from pathlib import PosixPath
 from typing import Optional, Tuple
 
+import torch
+from compressed_tensors.utils import force_cpu_offload
 from loguru import logger
 from torch.nn import Module
 from transformers import (
@@ -62,6 +64,16 @@ def pre_process(model_args: "ModelArguments"):
     # untie tie_word_embeddings weights
     patch_tied_tensors_bug(model_args.model)
 
+    # offload to cpu if possible
+    if "cuda" in str(model_args.oneshot_device) and torch.cuda.is_available():
+        # TODO: consider renaming function to something like "offload_dispatch_model"
+        # TODO: modify function to remove any hooks if they already exist (making sure
+        # to move to cpu when removing hook
+        force_cpu_offload(model_args.model, model_args.oneshot_device)
+
+    else:
+        logger.warning("CUDA is not available! Compressing model on CPU instead")
+
     # wrap model.save_pretrained
     modify_save_pretrained(model_args.model)
 

src/llmcompressor/pipelines/layer_sequential/pipeline.py

@@ -2,6 +2,7 @@
 
 import torch
 import tqdm
+from loguru import logger
 from torch.utils.data.dataloader import DataLoader
 
 from llmcompressor.core import LifecycleCallbacks, active_session
@@ -57,6 +58,16 @@ def __call__(
         """
         session = active_session()
 
+        # check for offloading
+        if model.device != torch.device("meta"):
+            logger.warning(
+                "Attemping to use sequential pipeline with a model which is not "
+                "offloaded to the cpu. Deploying a model in this way may lead to more "
+                "memory usage than is required. It is recommended to set "
+                '`oneshot_device="cuda"` or call `force_cpu_offload` on your model '
+                "before compressing"
+            )
+
         # find layers
         modifiers = session.get_modifiers()
         sequential_targets, _ = get_targets_from_modifiers(modifiers, model)

src/llmcompressor/pipelines/registry.py

@@ -18,6 +18,7 @@
 __all__ = ["CalibrationPipeline"]
 
 SEQUENTIAL_MODIFIERS = (AWQModifier, GPTQModifier, SparsityModifierBase)
+NEED_DATA = (SmoothQuantModifier, *SEQUENTIAL_MODIFIERS)
 
 
 class CalibrationPipeline(ABC, RegistryMixin):
@@ -60,7 +61,7 @@ def from_modifiers(
 
     @staticmethod
     def _validate_infer_pipeline(modifiers: List[Modifier]) -> str:
-        if any(isinstance(modifier, SEQUENTIAL_MODIFIERS) for modifier in modifiers):
+        if any(isinstance(modifier, NEED_DATA) for modifier in modifiers):
             return "sequential"
 
         active_qmods = _get_active_quant_modifiers(modifiers)
@@ -76,11 +77,6 @@ def _validate_infer_pipeline(modifiers: List[Modifier]) -> str:
             config = quant_modifier.resolve_quantization_config()
             if config.requires_calibration_data():
                 return "sequential"
-            else:
-                return "datafree"
-
-        if any(isinstance(modifier, SmoothQuantModifier) for modifier in modifiers):
-            return "sequential"
 
         return "datafree"
 

src/llmcompressor/pipelines/sequential/pipeline.py

@@ -2,6 +2,7 @@
 
 import torch
 from compressed_tensors.utils import get_execution_device
+from loguru import logger
 from torch.utils.data.dataloader import DataLoader
 from tqdm import tqdm
 
@@ -52,6 +53,16 @@ def __call__(
         """
         session = active_session()
 
+        # check for offloading
+        if model.device != torch.device("meta"):
+            logger.warning(
+                "Attemping to use sequential pipeline with a model which is not "
+                "offloaded to the cpu. Deploying a model in this way may lead to more "
+                "memory usage than is required. It is recommended to set "
+                '`oneshot_device="cuda"` or call `force_cpu_offload` on your model '
+                "before compressing"
+            )
+
         # prepare to trace subgraphs
         modifiers = session.get_modifiers()
         sequential_targets = get_targets_from_modifiers(modifiers, model)