fix initial device, update import

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

Author: kylesayrs

src/llmcompressor/pipelines/basic/pipeline.py

@@ -9,8 +9,7 @@
 from llmcompressor.modifiers.utils.pytorch_helpers import apply_pad_mask_to_batch
 from llmcompressor.pipelines.registry import CalibrationPipeline
 from llmcompressor.pytorch.utils.helpers import tensors_to_device
-from llmcompressor.utils.dev import dispatch_for_generation
-from llmcompressor.utils.helpers import calibration_forward_context
+from llmcompressor.utils import calibration_forward_context, dispatch_for_generation
 
 if TYPE_CHECKING:
     from llmcompressor.args.dataset_arguments import DatasetArguments

src/llmcompressor/pipelines/cache.py

@@ -61,7 +61,7 @@ def empty(cls, num_batches: int, offload_device: torch.device):
     def from_dataloader(
         cls,
         dataloader: torch.utils.data.DataLoader,
-        model_device: torch.device,
+        model_device: torch.device = torch.device("cpu"),
         mask_padding: bool = True,
         offload_device: Optional[torch.device] = torch.device("cpu"),
     ):

src/llmcompressor/pipelines/layer_sequential/helpers.py

@@ -6,7 +6,6 @@
 import torch
 import tqdm
 from compressed_tensors.quantization import find_name_or_class_matches
-from compressed_tensors.utils import get_execution_device
 from torch.nn import Module
 from torch.utils.data.dataloader import DataLoader
 
@@ -62,15 +61,14 @@ def capture_first_layer_intermediates(
     :param mask_padding: zero out padding tokens if True. This affects modifiers such as
         GPTQ and SparseGPT
     """
-    model_device = get_execution_device(model)
     intermediates = IntermediatesCache.empty(len(dataloader), torch.device("cpu"))
     signature = inspect.signature(first_layer.forward)
 
     with calibration_forward_context(model), early_stop_hook(first_layer):
         desc = "Preparing intermediates cache"
         for batch_index, batch in enumerate(tqdm.tqdm(dataloader, desc=desc)):
             batch = apply_pad_mask_to_batch(batch) if mask_padding else batch
-            batch = tensors_to_device(batch, model_device)
+            batch = tensors_to_device(batch, torch.device("cpu"))
 
             try:
                 model(**batch)

src/llmcompressor/pipelines/sequential/pipeline.py

@@ -1,7 +1,7 @@
 from typing import TYPE_CHECKING
 
 import torch
-from compressed_tensors.utils import disable_offloading, get_execution_device
+from compressed_tensors.utils import disable_offloading
 from torch.utils.data.dataloader import DataLoader
 from tqdm import tqdm
 
@@ -69,8 +69,7 @@ def __call__(
 
         with calibration_forward_context(model), DisableQuantization(model):
             # prepare intermediates cache
-            model_device = get_execution_device(model)
-            activations = IntermediatesCache.from_dataloader(dataloader, model_device)
+            activations = IntermediatesCache.from_dataloader(dataloader)
 
             for subgraph_index, subgraph in enumerate(subgraphs):
                 # prepare tqdm description texts