finish

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

Author: kylesayrs

src/llmcompressor/pipelines/registry.py

@@ -7,8 +7,9 @@
 from torch.utils.data.dataloader import DataLoader
 
 from llmcompressor.modifiers import Modifier
+from llmcompressor.modifiers.awq import AWQModifier
 from llmcompressor.modifiers.pruning.sparsegpt.sgpt_base import SparsityModifierBase
-from llmcompressor.modifiers.quantization import QuantizationModifier
+from llmcompressor.modifiers.quantization import GPTQModifier, QuantizationModifier
 from llmcompressor.modifiers.smoothquant import SmoothQuantModifier
 
 if TYPE_CHECKING:
@@ -58,11 +59,14 @@ def from_modifiers(
     @staticmethod
     def _infer_pipeline(modifiers: list[Modifier]) -> str:
         def _modifier_requires_calibration(modifier: Modifier):
-            if isinstance(modifier, QuantizationModifier):
+            if isinstance(
+                modifier,
+                (SmoothQuantModifier, SparsityModifierBase, GPTQModifier, AWQModifier),
+            ):
+                return True
+            elif isinstance(modifier, QuantizationModifier):
                 config = modifier.resolve_quantization_config()
                 return config.requires_calibration_data()
-            elif isinstance(modifier, (SmoothQuantModifier, SparsityModifierBase)):
-                return True
             else:
                 return False
 

tests/llmcompressor/pipelines/test_registry.py

@@ -1,28 +1,35 @@
 import pytest
 
-from llmcompressor.modifiers.quantization import QuantizationModifier, GPTQModifier
-from llmcompressor.modifiers.smoothquant import SmoothQuantModifier
 from llmcompressor.modifiers.awq import AWQModifier
-from llmcompressor.modifiers.pruning import SparseGPTQModifier, WandaPruningModifier
+from llmcompressor.modifiers.pruning import SparseGPTModifier, WandaPruningModifier
+from llmcompressor.modifiers.quantization import GPTQModifier, QuantizationModifier
+from llmcompressor.modifiers.smoothquant import SmoothQuantModifier
 from llmcompressor.modifiers.transform import QuIPModifier, SpinQuantModifier
-from llmcompressor.pipelines import CalibrationPipeline, SequentialPipeline, DataFreePipeline
+from llmcompressor.pipelines import (
+    CalibrationPipeline,
+    DataFreePipeline,
+    SequentialPipeline,
+)
 
 
-@pytest.mark.parametrize("modifiers", [
-    ([QuantizationModifier(scheme="FP8")], SequentialPipeline)
-    ([QuantizationModifier(scheme="W4A16")], DataFreePipeline)
-    ([GPTQModifier(scheme="FP8")], SequentialPipeline)
-    ([GPTQModifier(scheme="W4A16")], DataFreePipeline)
-    ([SmoothQuantModifier(), GPTQModifier(scheme="W4A16")], SequentialPipeline),
-    ([AWQModifier(scheme="W4A16")], SequentialPipeline)
-    ([AWQModifier(scheme="FP8")], SequentialPipeline)
-    ([SparseGPTQModifier()], SequentialPipeline)
-    ([WandaPruningModifier()], SequentialPipeline)
-    ([QuIPModifier()], DataFreePipeline)
-    ([SpinQuantModifier()], DataFreePipeline)
-    ([QuIPModifier(), QuantizationModifier(scheme="FP8")], SequentialPipeline)
-    ([QuIPModifier(), QuantizationModifier(scheme="W4A16")], DataFreePipeline)
-])
+@pytest.mark.parametrize(
+    "modifiers,exp_pipeline",
+    [
+        ([QuantizationModifier(scheme="FP8")], SequentialPipeline),
+        ([QuantizationModifier(scheme="W4A16")], DataFreePipeline),
+        ([GPTQModifier(scheme="FP8")], SequentialPipeline),
+        ([GPTQModifier(scheme="W4A16")], SequentialPipeline),
+        ([SmoothQuantModifier(), GPTQModifier(scheme="W4A16")], SequentialPipeline),
+        ([AWQModifier(scheme="W4A16")], SequentialPipeline),
+        ([AWQModifier(scheme="FP8")], SequentialPipeline),
+        ([SparseGPTModifier(sparsity=1.0)], SequentialPipeline),
+        ([WandaPruningModifier(sparsity=1.0)], SequentialPipeline),
+        ([QuIPModifier()], DataFreePipeline),
+        ([SpinQuantModifier()], DataFreePipeline),
+        ([QuIPModifier(), QuantizationModifier(scheme="FP8")], SequentialPipeline),
+        ([QuIPModifier(), QuantizationModifier(scheme="W4A16")], DataFreePipeline),
+    ],
+)
 def test_infer_pipeline(modifiers, exp_pipeline):
     pipeline = CalibrationPipeline.from_modifiers(modifiers)
-    assert isinstance(pipeline, exp_pipeline)
+    assert isinstance(pipeline, exp_pipeline)