Use align_module_device util (#1298)

## Purpose ##
* Standardization and clarity

## Changes ##
* Replace all uses of `_hf_hook.pre_forward` with `align_module_device`

## Testing ##
* `grep -r '_hf_hook.pre_forward' src/`

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

Author: kylesayrs

src/llmcompressor/modifiers/quantization/calibration.py

@@ -8,7 +8,7 @@
 )
 from compressed_tensors.quantization.lifecycle.forward import forward_quantize
 from compressed_tensors.quantization.utils import is_kv_cache_quant_scheme
-from compressed_tensors.utils.offload import is_module_offloaded, update_parameter_data
+from compressed_tensors.utils import align_module_device, update_parameter_data
 from loguru import logger
 from torch.nn import Module
 
@@ -72,27 +72,23 @@ def call_observer(module: Module, base_name: str, value: Optional[torch.Tensor]
     :param value: torch.Tensor to be passed to the observer for activations. If
         base_name is "weight", then the module's weight tensor will be used
     """
-    offloaded = is_module_offloaded(module)
-    if offloaded:
-        module._hf_hook.pre_forward(module)
-
-    if base_name == "weight":
-        value = module.weight
-        g_idx = getattr(module, "weight_g_idx", None)
-    elif value is not None:
-        g_idx = None
-    else:
-        raise ValueError("Must provide a value to observe if not using weight observer")
-
-    observer = getattr(module, f"{base_name}_observer")
-    updated_scale, updated_zero_point = observer(value, g_idx=g_idx)
+    with align_module_device(module):
+        if base_name == "weight":
+            value = module.weight
+            g_idx = getattr(module, "weight_g_idx", None)
+        elif value is not None:
+            g_idx = None
+        else:
+            raise ValueError(
+                "Must provide a value to observe if not using weight observer"
+            )
 
-    # update scale and zero point
-    update_parameter_data(module, updated_scale, f"{base_name}_scale")
-    update_parameter_data(module, updated_zero_point, f"{base_name}_zero_point")
+        observer = getattr(module, f"{base_name}_observer")
+        updated_scale, updated_zero_point = observer(value, g_idx=g_idx)
 
-    if offloaded:
-        module._hf_hook.post_forward(module, None)
+        # update scale and zero point
+        update_parameter_data(module, updated_scale, f"{base_name}_scale")
+        update_parameter_data(module, updated_zero_point, f"{base_name}_zero_point")
 
 
 def update_weight_zp_scale(module: Module):

src/llmcompressor/modifiers/smoothquant/base.py

@@ -2,7 +2,7 @@
 from typing import Callable, Dict, List, Optional, Tuple, Union
 
 import torch
-from compressed_tensors.utils.offload import is_module_offloaded
+from compressed_tensors.utils import align_module_device
 from loguru import logger
 from pydantic import Field
 from torch.nn import Module
@@ -290,22 +290,16 @@ def _apply_smoothing(self, model: Module):
 
             @torch.no_grad()
             def smooth(module):
-                offloaded = is_module_offloaded(module)
-                if offloaded:
-                    module._hf_hook.pre_forward(module)
-
-                if module in balance_layers:
-                    module.weight.mul_(scales.view(1, -1))
-                elif module == smooth_layer:
-                    if module.weight.ndim == 1:
-                        module.weight.div_(scales)
-                    else:
-                        module.weight.div_(scales.view(-1, 1))
-                    if hasattr(module, "bias") and module.bias is not None:
-                        module.bias.div_(scales)
-
-                if offloaded:
-                    module._hf_hook.post_forward(module, None)
+                with align_module_device(module):
+                    if module in balance_layers:
+                        module.weight.mul_(scales.view(1, -1))
+                    elif module == smooth_layer:
+                        if module.weight.ndim == 1:
+                            module.weight.div_(scales)
+                        else:
+                            module.weight.div_(scales.view(-1, 1))
+                        if hasattr(module, "bias") and module.bias is not None:
+                            module.bias.div_(scales)
 
             parent = get_fsdp_parent(mapping.smooth_name, model)
             if parent is not None:
@@ -330,15 +324,9 @@ def _calculate_smoothing_scales(
         # get the channel-wise dynamic range for each layer to be balanced
         weight_scales = []
         for layer in balance_layers:
-            offloaded = is_module_offloaded(layer)
-            if offloaded:
-                layer._hf_hook.pre_forward(layer)
-
-            scale = layer.weight.abs().max(dim=0, keepdim=True)[0]
-            weight_scales.append(scale)
-
-            if offloaded:
-                layer._hf_hook.post_forward(layer, None)
+            with align_module_device(layer):
+                scale = layer.weight.abs().max(dim=0, keepdim=True)[0]
+                weight_scales.append(scale)
 
         weight_scales = 2.0 * torch.cat(weight_scales, dim=0).max(dim=0)[0]
 

src/llmcompressor/transformers/compression/helpers.py

@@ -5,8 +5,8 @@
 import torch
 from accelerate import infer_auto_device_map, init_empty_weights
 from accelerate.accelerator import get_state_dict_offloaded_model
-from compressed_tensors import is_module_offloaded
 from compressed_tensors.quantization.utils import iter_named_leaf_modules, module_type
+from compressed_tensors.utils import align_module_device
 from torch.nn.modules import Linear
 from tqdm import tqdm
 from transformers import AutoModelForCausalLM
@@ -298,18 +298,14 @@ def is_sparse_compression_target(
     :return: whether or not the module is a target for sparsity compression,
         i.e True if it is sparse and follows the sparsity structure, else False
     """
-    offloaded = is_module_offloaded(module)
-    if offloaded:
-        module._hf_hook.pre_forward(module)
-
-    result = (
-        hasattr(module, "weight")
-        and tensor_sparsity(module.weight) >= sparsity_threshold
-        and tensor_follows_mask_structure(tensor=module.weight, mask=sparsity_structure)
-    )
-
-    if offloaded:
-        module._hf_hook.post_forward(module, None)
+    with align_module_device(module):
+        result = (
+            hasattr(module, "weight")
+            and tensor_sparsity(module.weight) >= sparsity_threshold
+            and tensor_follows_mask_structure(
+                tensor=module.weight, mask=sparsity_structure
+            )
+        )
 
     return result