Modernized typehints in core/lifecycle.py

siddhaka · siddhaka · commit 0a08e487a4ea · 2025-10-14T19:33:08.000-05:00
Signed-off-by: siddhaka &lt;siddhaka@gmail.com&gt;
diff --git a/examples/awq/qwen3_coder_moe_example.py b/examples/awq/qwen3_coder_moe_example.py
@@ -31,7 +31,7 @@
 def get_calib_dataset(tokenizer):
     ds = load_dataset(
         DATASET_ID,
-        split=f"{DATASET_SPLIT}[:{NUM_CALIBRATION_SAMPLES*10}]",
+        split=f"{DATASET_SPLIT}[:{NUM_CALIBRATION_SAMPLES * 10}]",
     )
 
     def preprocess(example):
diff --git a/src/llmcompressor/modeling/fuse.py b/src/llmcompressor/modeling/fuse.py
@@ -47,8 +47,9 @@ def fuse_norm_linears(norm: torch.nn.Module, linears: Iterable[torch.nn.Linear])
     for linear in linears:
         # NOTE: spinquant does this op in float64
         exec_device = get_execution_device(norm)
-        with align_module_device(norm, exec_device), align_module_device(
-            linear, exec_device
+        with (
+            align_module_device(norm, exec_device),
+            align_module_device(linear, exec_device),
         ):
             weight_dtype = linear.weight.dtype
             new_weight = linear.weight.to(PRECISION) * norm.weight.to(PRECISION)
diff --git a/src/llmcompressor/modifiers/awq/base.py b/src/llmcompressor/modifiers/awq/base.py
@@ -155,9 +155,9 @@ def validate_awq_after(model: "AWQModifier") -> "AWQModifier":
             for group in config.config_groups.values()
             if group.weights is not None
         )
-        assert (
-            len(num_bits_set) == 1
-        ), "In AWQ, all config groups must use the same configuration for num_bits"
+        assert len(num_bits_set) == 1, (
+            "In AWQ, all config groups must use the same configuration for num_bits"
+        )
 
         model._num_bits = next(iter(num_bits_set))
 
@@ -166,9 +166,9 @@ def validate_awq_after(model: "AWQModifier") -> "AWQModifier":
             for group in config.config_groups.values()
             if group.weights is not None
         )
-        assert (
-            len(symmetric_set) == 1
-        ), "In AWQ, all config groups must use the same configuration for symmetric"
+        assert len(symmetric_set) == 1, (
+            "In AWQ, all config groups must use the same configuration for symmetric"
+        )
 
         model._symmetric = next(iter(symmetric_set))
 
@@ -177,9 +177,9 @@ def validate_awq_after(model: "AWQModifier") -> "AWQModifier":
             for group in config.config_groups.values()
             if group.weights is not None
         )
-        assert (
-            len(group_size_set) == 1
-        ), "In AWQ, all config groups must use the same configuration for group_size"
+        assert len(group_size_set) == 1, (
+            "In AWQ, all config groups must use the same configuration for group_size"
+        )
 
         model._group_size = next(iter(group_size_set))
 
@@ -316,7 +316,7 @@ def _set_resolved_mappings(self, model: Module) -> None:
                 )
             ):
                 pbar.set_description(
-                    f"Resolving mapping {mapping_idx+1}/{len(self.mappings)}"
+                    f"Resolving mapping {mapping_idx + 1}/{len(self.mappings)}"
                     f" ({num_skipped_mappings} skipped)"
                 )
 
@@ -452,9 +452,11 @@ def _apply_smoothing(self, model: Module) -> None:
             balance_layers = mapping.balance_layers
             parent_module = mapping.parent
 
-            with align_modules(
-                [parent_module, smooth_layer, *balance_layers]
-            ), calibration_forward_context(model), HooksMixin.disable_hooks():
+            with (
+                align_modules([parent_module, smooth_layer, *balance_layers]),
+                calibration_forward_context(model),
+                HooksMixin.disable_hooks(),
+            ):
                 # [STEP 1]: Compute per-channel mean of normalised weights
                 # All layer weights are concatted together
                 weight = torch.cat([bl.weight for bl in balance_layers], dim=0)
@@ -653,9 +655,9 @@ def _compute_best_scale(
                 "https://github.com/vllm-project/llm-compressor/issues"
             )
 
-        assert (
-            torch.isnan(best_scales).sum() == 0
-        ), f"Nan found in scales: {best_scales}"
+        assert torch.isnan(best_scales).sum() == 0, (
+            f"Nan found in scales: {best_scales}"
+        )
 
         return best_scales.detach().cpu()
 
diff --git a/src/llmcompressor/modifiers/pruning/sparsegpt/base.py b/src/llmcompressor/modifiers/pruning/sparsegpt/base.py
@@ -123,9 +123,11 @@ def compress_modules(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,
diff --git a/src/llmcompressor/modifiers/pruning/utils/pytorch/layer_mask.py b/src/llmcompressor/modifiers/pruning/utils/pytorch/layer_mask.py
@@ -23,7 +23,7 @@ def setup_mask_for_param(param: Parameter, mask: torch.Tensor) -> torch.Tensor:
 
     if mask.shape != param.data.shape:
         raise ValueError(
-            f"Mask shape {mask.shape} does not match " f"param shape {param.data.shape}"
+            f"Mask shape {mask.shape} does not match param shape {param.data.shape}"
         )
 
     if mask.dtype != torch.bool:
diff --git a/src/llmcompressor/modifiers/pruning/wanda/base.py b/src/llmcompressor/modifiers/pruning/wanda/base.py
@@ -108,8 +108,10 @@ def compress_modules(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,
diff --git a/src/llmcompressor/modifiers/quantization/gptq/base.py b/src/llmcompressor/modifiers/quantization/gptq/base.py
@@ -249,11 +249,12 @@ def compress_modules(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,
diff --git a/src/llmcompressor/recipe/utils.py b/src/llmcompressor/recipe/utils.py
@@ -48,9 +48,7 @@ def _parse_recipe_from_md(file_path, yaml_str):
     else:
         # fail if we know whe should have extracted front matter out
         raise RuntimeError(
-            "Could not extract YAML front matter from recipe card:" " {}".format(
-                file_path
-            )
+            "Could not extract YAML front matter from recipe card: {}".format(file_path)
         )
     return yaml_str
 
diff --git a/src/llmcompressor/utils/dev.py b/src/llmcompressor/utils/dev.py
@@ -70,9 +70,12 @@ def patched(cls, *args, **kwargs):
 
         return model
 
-    with tempfile.TemporaryDirectory() as tmp_dir, patch_attr(
-        model_class, "from_pretrained", patched
-    ), skip_weights_initialize(), patch_transformers_logger_level():
+    with (
+        tempfile.TemporaryDirectory() as tmp_dir,
+        patch_attr(model_class, "from_pretrained", patched),
+        skip_weights_initialize(),
+        patch_transformers_logger_level(),
+    ):
         yield
 
 
diff --git a/src/llmcompressor/utils/helpers.py b/src/llmcompressor/utils/helpers.py
@@ -1049,8 +1049,11 @@ def calibration_forward_context(model: torch.nn.Module):
     - Disable train mode and enable eval mode
     - Disable hf kernels which could bypass hooks
     """
-    with torch.no_grad(), disable_cache(model), eval_context(model), disable_hf_kernels(
-        model
+    with (
+        torch.no_grad(),
+        disable_cache(model),
+        eval_context(model),
+        disable_hf_kernels(model),
     ):
         yield
 
diff --git a/src/llmcompressor/utils/metric_logging.py b/src/llmcompressor/utils/metric_logging.py
@@ -144,10 +144,7 @@ def __exit__(self, _exc_type, _exc_val, _exc_tb):
                 total_memory = int(gpu_usage[i][1])  # GB
                 patch.log(
                     "METRIC",
-                    (
-                        f"GPU {i} | usage: {perc:.2f}%"
-                        f" | total memory: {total_memory} GB"
-                    ),
+                    (f"GPU {i} | usage: {perc:.2f}% | total memory: {total_memory} GB"),
                 )
 
         compressed_size = get_layer_size_mb(self.module)
diff --git a/tests/examples/test_example_scripts.py b/tests/examples/test_example_scripts.py
@@ -62,9 +62,9 @@ def verify_2of4_w4a16_output(tmp_path: Path, example_dir: str):
 def verify_w4a4_fp4_output(tmp_path: Path, example_dir: str):
     # verify the expected directory was generated
     nvfp4_dirs: List[Path] = [p for p in tmp_path.rglob("*-NVFP4") if p.is_dir()]
-    assert (
-        len(nvfp4_dirs)
-    ) == 1, f"did not find exactly one generated folder: {nvfp4_dirs}"
+    assert (len(nvfp4_dirs)) == 1, (
+        f"did not find exactly one generated folder: {nvfp4_dirs}"
+    )
 
     # verify the format in the generated config
     config_json = json.loads((nvfp4_dirs[0] / "config.json").read_text())
diff --git a/tests/llmcompressor/modeling/test_calib_qwen3.py b/tests/llmcompressor/modeling/test_calib_qwen3.py
@@ -59,9 +59,9 @@ def hook_fn(i, module, input, output):
             _ = moe_layer(sample)
 
         # Assert all experts are used
-        assert all(
-            expert_triggered
-        ), f"Not all experts were triggered: {expert_triggered}"
+        assert all(expert_triggered), (
+            f"Not all experts were triggered: {expert_triggered}"
+        )
 
 
 @requires_gpu
diff --git a/tests/llmcompressor/modifiers/logarithmic_equalization/test_base.py b/tests/llmcompressor/modifiers/logarithmic_equalization/test_base.py
@@ -20,9 +20,9 @@ def test_logarithmic_equalization_is_registered():
         mappings=mappings,
     )
 
-    assert isinstance(
-        modifier, LogarithmicEqualizationModifier
-    ), "PyTorch LogarithmicEqualizationModifier not registered"
+    assert isinstance(modifier, LogarithmicEqualizationModifier), (
+        "PyTorch LogarithmicEqualizationModifier not registered"
+    )
     assert isinstance(modifier, SmoothQuantModifier)
     assert modifier.smoothing_strength == smoothing_strength
     assert modifier.mappings == mappings
diff --git a/tests/llmcompressor/modifiers/pruning/sparsegpt/test_base.py b/tests/llmcompressor/modifiers/pruning/sparsegpt/test_base.py
@@ -17,6 +17,6 @@ def test_sparse_gpt_is_registered():
         targets=targets,
     )
 
-    assert isinstance(
-        type_, SparseGPTModifier
-    ), "PyTorch SparseGPTModifier not registered"
+    assert isinstance(type_, SparseGPTModifier), (
+        "PyTorch SparseGPTModifier not registered"
+    )
diff --git a/tests/llmcompressor/modifiers/pruning/wanda/test_base.py b/tests/llmcompressor/modifiers/pruning/wanda/test_base.py
@@ -17,6 +17,6 @@ def test_wanda_is_registered():
         targets=targets,
     )
 
-    assert isinstance(
-        type_, WandaPruningModifier
-    ), "PyTorch WandaPruningModifier not registered"
+    assert isinstance(type_, WandaPruningModifier), (
+        "PyTorch WandaPruningModifier not registered"
+    )
diff --git a/tests/llmcompressor/modifiers/smoothquant/test_base.py b/tests/llmcompressor/modifiers/smoothquant/test_base.py
@@ -17,9 +17,9 @@ def test_smooth_quant_is_registered():
         mappings=mappings,
     )
 
-    assert isinstance(
-        modifier, SmoothQuantModifier
-    ), "PyTorch SmoothQuant not registered"
+    assert isinstance(modifier, SmoothQuantModifier), (
+        "PyTorch SmoothQuant not registered"
+    )
     assert modifier.smoothing_strength == smoothing_strength
     assert modifier.mappings == mappings
 
diff --git a/tests/llmcompressor/modifiers/utils/test_hooks.py b/tests/llmcompressor/modifiers/utils/test_hooks.py
@@ -139,8 +139,9 @@ 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
diff --git a/tests/llmcompressor/pytorch/modifiers/pruning/constant/test_pytorch.py b/tests/llmcompressor/pytorch/modifiers/pruning/constant/test_pytorch.py
@@ -154,6 +154,6 @@ def test_constant_pruning_pytorch_is_registered():
         **kwargs,
     )
 
-    assert isinstance(
-        type_, ConstantPruningModifier
-    ), "PyTorch ConstantPruningModifier not registered"
+    assert isinstance(type_, ConstantPruningModifier), (
+        "PyTorch ConstantPruningModifier not registered"
+    )
diff --git a/tests/llmcompressor/pytorch/modifiers/pruning/wanda/test_pytorch.py b/tests/llmcompressor/pytorch/modifiers/pruning/wanda/test_pytorch.py
@@ -18,6 +18,6 @@ def test_wanda_pytorch_is_registered():
         targets=targets,
     )
 
-    assert isinstance(
-        type_, WandaPruningModifier
-    ), "PyTorch ConstantPruningModifier not registered"
+    assert isinstance(type_, WandaPruningModifier), (
+        "PyTorch ConstantPruningModifier not registered"
+    )
diff --git a/tests/llmcompressor/transformers/compression/test_compress_tensor_utils.py b/tests/llmcompressor/transformers/compression/test_compress_tensor_utils.py
@@ -186,12 +186,12 @@ def test_quant_model_reload(format, dtype, tmp_path):
 
     for name, module in model.named_modules():
         if hasattr(module, "quantization_scheme"):
-            assert (
-                module.weight.dtype == dtype
-            ), f"Module {name} has incorrect weight dtype"
-            assert (
-                module.quantization_status == QuantizationStatus.FROZEN
-            ), f"Module {name} has incorrect quantization status"
+            assert module.weight.dtype == dtype, (
+                f"Module {name} has incorrect weight dtype"
+            )
+            assert module.quantization_status == QuantizationStatus.FROZEN, (
+                f"Module {name} has incorrect quantization status"
+            )
 
     # Save to disk
     model.save_pretrained(
@@ -373,14 +373,14 @@ def test_compressor_stacking(model_stub, recipe, sparse_format, quant_format, tm
         model, sparsity_config_or_format=sparse_format, quantization_format=quant_format
     )
 
-    assert (
-        compressor.sparsity_compressor is not None
-    ), "Sparse compressor not initialized"
+    assert compressor.sparsity_compressor is not None, (
+        "Sparse compressor not initialized"
+    )
     assert compressor.sparsity_config.format == sparse_format
 
-    assert (
-        compressor.quantization_compressor is not None
-    ), "Quantization compressor not initialized"
+    assert compressor.quantization_compressor is not None, (
+        "Quantization compressor not initialized"
+    )
 
     compressor.compress_model(model)
     compressor.decompress_model(model)
@@ -459,9 +459,9 @@ def test_sparse_24_compressor_is_lossless(model_stub, recipe, sparse_format, tmp
     # the model instead
     compressor = ModelCompressor.from_compression_config(compression_config)
 
-    assert (
-        compressor.sparsity_compressor is not None
-    ), "Sparse compressor not initialized"
+    assert compressor.sparsity_compressor is not None, (
+        "Sparse compressor not initialized"
+    )
     assert compressor.sparsity_config.format == sparse_format
 
     compressor.decompress(model_path=path, model=empty_model)
diff --git a/tests/llmcompressor/transformers/compression/test_recipe_parsing.py b/tests/llmcompressor/transformers/compression/test_recipe_parsing.py
@@ -90,8 +90,8 @@ def test_oneshot_accepts_multiple_recipe_formats(setup_model_and_config, recipe)
     )
 
     output_path = Path(setup_model_and_config["output_dir"])
-    assert output_path.exists() and any(
-        output_path.iterdir()
-    ), f"No output artifacts found in: {output_path}"
+    assert output_path.exists() and any(output_path.iterdir()), (
+        f"No output artifacts found in: {output_path}"
+    )
 
     reset_session()
diff --git a/tests/llmcompressor/transformers/sparsegpt/test_consecutive_runs.py b/tests/llmcompressor/transformers/sparsegpt/test_consecutive_runs.py
@@ -38,9 +38,9 @@ def test_consecutive_runs_small(config, tmp_path):
 @pytest.mark.integration
 @pytest.mark.parametrize("config", parse_params(GPU_CONFIGS_DIRECTORY))
 def test_consecutive_runs_gpu(config, tmp_path):
-    assert not is_model_ct_quantized_from_path(
-        config["model"]
-    ), "The provided model is quantized. Please use a dense model."
+    assert not is_model_ct_quantized_from_path(config["model"]), (
+        "The provided model is quantized. Please use a dense model."
+    )
     model = AutoModelForCausalLM.from_pretrained(
         config["model"], device_map=config["device"], torch_dtype="auto"
     )
diff --git a/tests/testing_utils.py b/tests/testing_utils.py
@@ -128,9 +128,9 @@ def parse_params(configs_directory: Union[list, str]) -> List[dict]:
     config_dicts = []
 
     def _parse_configs_dir(current_config_dir):
-        assert os.path.isdir(
-            current_config_dir
-        ), f"Config_directory {current_config_dir} is not a directory"
+        assert os.path.isdir(current_config_dir), (
+            f"Config_directory {current_config_dir} is not a directory"
+        )
 
         for file in os.listdir(current_config_dir):
             config_path = os.path.join(current_config_dir, file)

Original file line number	Diff line number	Diff line change
`@@ -31,7 +31,7 @@`
`31`	`31`	`def get_calib_dataset(tokenizer):`
`32`	`32`	`ds = load_dataset(`
`33`	`33`	`DATASET_ID,`
`34`		`- split=f"{DATASET_SPLIT}[:{NUM_CALIBRATION_SAMPLES*10}]",`
	`34`	`+ split=f"{DATASET_SPLIT}[:{NUM_CALIBRATION_SAMPLES * 10}]",`
`35`	`35`	`)`
`36`	`36`
`37`	`37`	`def preprocess(example):`
Original file line number	Diff line number	Diff line change
`@@ -23,7 +23,7 @@ def setup_mask_for_param(param: Parameter, mask: torch.Tensor) -> torch.Tensor:`
`23`	`23`
`24`	`24`	`if mask.shape != param.data.shape:`
`25`	`25`	`raise ValueError(`
`26`		`- f"Mask shape {mask.shape} does not match " f"param shape {param.data.shape}"`
	`26`	`+ f"Mask shape {mask.shape} does not match param shape {param.data.shape}"`
`27`	`27`	`)`
`28`	`28`
`29`	`29`	`if mask.dtype != torch.bool:`