Fix Multi-Context Manager Syntax for Python 3.9 Compatibility (#1313)

Brings back #1287 since it was reverted (as full tests did not run)


Changes:

This PR resolves [Issue
#1250](#1250),
where the codebase relied on Python 3.10+ syntax for multiple context
managers in a single `with` statement (using parentheses).


 **Refactored Affected Modules**
- Updated `src/llmcompressor/modifiers/obcq/base.py`,
`src/llmcompressor/modifiers/pruning/wanda/base.py`,
`src/llmcompressor/modifiers/quantization/gptq/base.py`,
`src/llmcompressor/pipelines/sequential/helpers.py`, and
`src/llmcompressor/utils/helpers.py`
   - Modified `tests/llmcompressor/modifiers/utils/test_hooks.py` 

### Testing
- Verified that all modified modules run correctly under Python 3.9 and
3.10.
- Ran the updated test suite (`test_hooks.py`) to confirm hook disabling
behavior remains intact.
- Confirmed no regressions in sparsification, quantization, or tracing
pipelines.

### Related Issues
- Fixes
[#1250](#1250)

---------

Signed-off-by: Rahul Tuli <[email protected]>
Co-authored-by: Dipika Sikka <[email protected]>

Author: rahul-tuli

src/llmcompressor/modifiers/obcq/base.py

@@ -119,11 +119,9 @@ def on_sequential_batch_end(self):
             num_samples = self._num_samples[module]
 
             logger.info(f"Sparsifying {name} using {num_samples} samples")
-            with (
-                torch.no_grad(),
-                align_module_device(module),
-                CompressionLogger(module) as comp_logger,
-            ):
+            with torch.no_grad(), align_module_device(module), CompressionLogger(
+                module
+            ) as comp_logger:
                 loss, sparsified_weight = sparsify_weight(
                     module=module,
                     hessians_dict=self._hessians,

src/llmcompressor/modifiers/pruning/wanda/base.py

@@ -103,10 +103,8 @@ def on_sequential_batch_end(self):
             num_samples = self._num_samples[module]
 
             logger.info(f"Sparsifying {name} using {num_samples} samples")
-            with (
-                torch.no_grad(),
-                align_module_device(module),
-                CompressionLogger(module),
+            with torch.no_grad(), align_module_device(module), CompressionLogger(
+                module
             ):
                 sparsified_weight = sparsify_weight(
                     module=module,

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

@@ -337,12 +337,11 @@ def on_sequential_batch_end(self):
             quant_args = getattr_chain(module, "quantization_scheme.weights")
 
             logger.info(f"Quantizing {name} using {num_samples} samples")
-            with (
-                torch.no_grad(),
-                align_module_device(module),
-                self._maybe_onload_hessian(module),
-                CompressionLogger(module) as comp_logger,
-            ):
+            with torch.no_grad(), align_module_device(
+                module
+            ), self._maybe_onload_hessian(module), CompressionLogger(
+                module
+            ) as comp_logger:
                 loss, quantized_weight, scale, zero_point, g_idx = quantize_weight(
                     module=module,
                     quant_args=quant_args,

src/llmcompressor/pipelines/sequential/helpers.py

@@ -71,10 +71,7 @@ def trace_subgraphs(
     concrete_args = populate_concrete_args(model, sample_input)
 
     # trace
-    with (
-        calibration_forward_context(model),
-        HooksMixin.disable_hooks(),
-    ):
+    with calibration_forward_context(model), HooksMixin.disable_hooks():
         graph = GraphModule(
             model,
             tracer.trace(

tests/llmcompressor/modifiers/utils/test_hooks.py

@@ -139,9 +139,8 @@ def test_disable_hooks_composable():
     handle_b = mod_b.register_hook(model.linear2, mod_b.hook, "forward_pre")
 
     # composing two keeps
-    with (
-        HooksMixin.disable_hooks(keep=set([handle_b])),
-        HooksMixin.disable_hooks(keep=set([handle_a])),
+    with HooksMixin.disable_hooks(keep=set([handle_b])), HooksMixin.disable_hooks(
+        keep=set([handle_a])
     ):
         model(model.dummy_inputs)
     assert mod_a.hook_called and mod_b.hook_called