SQ and QM: Remove torch.cuda.empty_cache, use calibration_forward_context

src/llmcompressor/modifiers/quantization/quantization/base.py

@@ -32,6 +32,7 @@
     run_calibration_forward,
 )
 from llmcompressor.observers.helpers import get_observer_token_count
+from llmcompressor.utils.helpers import calibration_forward_context
 
 __all__ = ["QuantizationModifier"]
 
@@ -309,18 +310,13 @@ def _calibrate(self, module: Module):
             f"{len(self.calibration_dataloader_)} samples..."
         )
 
-        module_training = module.training
-        module.eval()
-
-        run_calibration_forward(
-            module,
-            self.calibration_dataloader_,
-            self.num_calibration_steps,
-            self.calibration_function_,
-        )
-
-        if module_training:
-            module.train()
+        with calibration_forward_context(module):
+            run_calibration_forward(
+                module,
+                self.calibration_dataloader_,
+                self.num_calibration_steps,
+                self.calibration_function_,
+            )
 
     def _check_token_distribution(
         self, model: Module, threshold: Optional[float] = None

src/llmcompressor/modifiers/smoothquant/base.py

@@ -14,6 +14,7 @@
 )
 from llmcompressor.modifiers.utils.pytorch_helpers import run_calibration_forward
 from llmcompressor.utils.fsdp.helpers import get_fsdp_parent
+from llmcompressor.utils.helpers import calibration_forward_context
 from llmcompressor.utils.pytorch.module import (
     get_layers,
     get_matching_layer,
@@ -250,12 +251,13 @@ def _calibrate(self, model: Module, calibration_dataloader: List):
                 " CompressionSession to run the SmoothQuant modifier"
             )
 
-        run_calibration_forward(
-            model,
-            calibration_dataloader,
-            self.num_calibration_steps,
-            self.calibration_function,
-        )
+        with calibration_forward_context(model):
+            run_calibration_forward(
+                model,
+                calibration_dataloader,
+                self.num_calibration_steps,
+                self.calibration_function,
+            )
 
         # remove the hooks now that we are done calibrating
         self.remove_hooks()
@@ -313,9 +315,6 @@ def smooth(module):
                     smooth(layer)
                 smooth(smooth_layer)
 
-        # clear out allocated smoothing scales
-        torch.cuda.empty_cache()
-
     def _calculate_smoothing_scales(
         self, balance_layers: List[Module], activation_scales: torch.Tensor
     ) -> List[float]:

src/llmcompressor/modifiers/utils/pytorch_helpers.py

@@ -81,10 +81,6 @@ def run_calibration_forward(
         with torch.no_grad():
             forward_fn(batch, module=model)
 
-        # TODO: not ideal, figure out where we aren't freeing memory instead
-        # currently without this we run OOM on the 2nd forward pass
-        torch.cuda.empty_cache()
-
 
 def is_moe_model(model: Module) -> bool:
     """

tests/llmcompressor/transformers/sparsification/test_compress_tensor_utils.py

@@ -672,7 +672,7 @@ def test_correct_compressor_inferred(
     if is_24:
         weights = _make_24_sparse(weights)
     else:
-        weights[0, :] = torch.ones(4, )  # guarantee not 24 sparse
+        weights[0, :] = torch.ones((4,))  # guarantee not 24 sparse
 
     quantization_config = _quantization_config_from_string(quant_style, quant_type)
     quantization_args = quantization_config.config_groups["group_0"].weights