use calib forward context

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

Author: kylesayrs

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]: