sq: remove empty_cache, add calib context

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

Author: kylesayrs

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

@@ -106,10 +106,6 @@ def run_calibration_forward(
                 # move on to next calibration sample
                 intermediates.append((e.args, e.kwargs))
 
-        # 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()
-
     return intermediates