[Actorder] Fix GPTQ actorder logic, only apply actorder to weight group args (#1815)

SUMMARY:
- Don't set actorder to static if running channel quantization
- The modifier level actoder value is also getting serialized
incorrectly - we set it to None if the strategy is not GROUP but still
missing a step to serialize it correctlt when it is None
- FIx test case which is using an incorrect activation quantization
strategy
-  Update compress / decompress test case 

Testing
- Fixes failing tests
- There is still a bug where `Sentinel` values are not serialized
correctly

---------

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

Author: dsikka

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

@@ -3,7 +3,11 @@
 from typing import Dict, List, Optional, Tuple, Union
 
 import torch
-from compressed_tensors.quantization import QuantizationConfig, QuantizationScheme
+from compressed_tensors.quantization import (
+    QuantizationConfig,
+    QuantizationScheme,
+    QuantizationStrategy,
+)
 from compressed_tensors.quantization.quant_args import ActivationOrdering
 from compressed_tensors.utils import (
     align_module_device,
@@ -107,6 +111,7 @@ class GPTQModifier(Modifier, QuantizationMixin):
     sequential_targets: Union[str, List[str], None] = None
     block_size: int = 128
     dampening_frac: Optional[float] = 0.01
+    # TODO: this does not serialize / will be incorrectly written
     actorder: Optional[Union[ActivationOrdering, Sentinel]] = Sentinel("static")
     offload_hessians: bool = False
 
@@ -149,9 +154,11 @@ def resolve_actorder(existing):
 
         for scheme in config.config_groups.values():
             assert isinstance(scheme, QuantizationScheme)
-            if scheme.weights is not None:
+            if (
+                getattr_chain(scheme, "weights.strategy", None)
+                == QuantizationStrategy.GROUP
+            ):
                 scheme.weights.actorder = resolve_actorder(scheme.weights.actorder)
-
         return config
 
     def on_initialize(self, state: State, **kwargs) -> bool:

tests/llmcompressor/modifiers/quantization/test_base.py

@@ -1,6 +1,7 @@
 from contextlib import nullcontext
 
 import pytest
+from compressed_tensors.quantization import QuantizationArgs, QuantizationScheme
 
 from llmcompressor.modifiers.quantization import GPTQModifier
 
@@ -107,3 +108,37 @@ def test_actorder_resolution(
         assert resolved.config_groups["group_0"].weights.actorder == expected_0
         assert resolved.config_groups["group_1"].input_activations.actorder is None
         assert resolved.config_groups["group_1"].weights.actorder == expected_1
+
+
+@pytest.mark.parametrize(
+    "strategies,actorder",
+    [
+        (["group"], None),
+        (["group"], "static"),
+        (["group"], "group"),
+        (["channel", "group"], None),
+        (["channel", "group"], "static"),
+        (["channel", "group"], "group"),
+        (["group", "channel"], None),
+        (["group", "channel"], "static"),
+        (["group", "channel"], "group"),
+    ],
+)
+def test_config_resolution(strategies, actorder):
+    config_groups = {
+        str(index): QuantizationScheme(
+            targets=[],
+            weights=QuantizationArgs(
+                strategy=strategy, group_size=(128 if strategy == "group" else None)
+            ),
+        )
+        for index, strategy in enumerate(strategies)
+    }
+
+    modifier = GPTQModifier(config_groups=config_groups, actorder=actorder)
+    modifier.resolve_quantization_config()
+
+    # validate that actorder was applied
+    for config_group in modifier.config_groups.values():
+        if config_group.weights.strategy == "group":
+            assert config_group.weights.actorder == actorder

tests/llmcompressor/transformers/sparsification/test_compress_tensor_utils.py

@@ -348,7 +348,6 @@ def test_compressor_stacking(model_stub, recipe, sparse_format, quant_format, tm
     concatenate_data = False
     num_calibration_samples = 64
     splits = {"calibration": "train[:10%]"}
-    empty_model = AutoModelForCausalLM.from_pretrained(model_stub, torch_dtype="auto")
 
     oneshot(
         model=model_stub,
@@ -357,29 +356,18 @@ def test_compressor_stacking(model_stub, recipe, sparse_format, quant_format, tm
         recipe=recipe,
         concatenate_data=concatenate_data,
         splits=splits,
-        clear_sparse_session=False,
     )
 
     # Fetch the oneshot model
     model = get_session_model()
     og_state_dict = model.state_dict()
     path = tmp_path / "compressed"
 
-    # Compress and save
-    model.save_pretrained(
-        path,
-        quantization_format=quant_format,
-        save_compressed=True,
-    )
-
-    # Verify config on disk
-    config = AutoConfig.from_pretrained(path)
-    compression_config = getattr(config, QUANTIZATION_CONFIG_NAME, None)
-    quant_config = ModelCompressor.parse_quantization_config(compression_config)
-
     # As HFQuantizer doesn't decompress the model, use the compressor to decompress
     # the model instead
-    compressor = ModelCompressor.from_compression_config(compression_config)
+    compressor = ModelCompressor.from_pretrained_model(
+        model, sparsity_config=sparse_format, quantization_format=quant_format
+    )
 
     assert (
         compressor.sparsity_compressor is not None
@@ -389,16 +377,15 @@ def test_compressor_stacking(model_stub, recipe, sparse_format, quant_format, tm
     assert (
         compressor.quantization_compressor is not None
     ), "Quantization compressor not initialized"
-    assert quant_config["format"] == quant_format
 
+    compressor.compress_model(model)
+    compressor.decompress_model(model)
     compressor.quantization_config.quantization_status = QuantizationStatus.FROZEN
-    compressor.decompress(model_path=path, model=empty_model)
 
     # Verify the abs difference between the decompressed model
     # and the original model
-    reconstructed_state_dict = empty_model.state_dict()
-    assert len(og_state_dict) == len(reconstructed_state_dict)
-    for key in og_state_dict.keys():
+    reconstructed_state_dict = model.state_dict()
+    for key in reconstructed_state_dict.keys():
         dense_tensor = og_state_dict[key].to(device)
         reconstructed_tensor = reconstructed_state_dict[key].to(device)
         assert dense_tensor.dtype == reconstructed_tensor.dtype
@@ -409,6 +396,16 @@ def test_compressor_stacking(model_stub, recipe, sparse_format, quant_format, tm
             assert not torch.any(diff > 0.025), f"Max diff: {torch.max(diff)}"
         else:
             assert torch.equal(dense_tensor, reconstructed_tensor)
+
+    # Recompress and save; validate correct formats used
+    model.save_pretrained(path)
+    config = AutoConfig.from_pretrained(path)
+    compression_config = getattr(config, QUANTIZATION_CONFIG_NAME, None)
+    quant_config = ModelCompressor.parse_quantization_config(compression_config)
+    sparsity_config = ModelCompressor.parse_sparsity_config(compression_config)
+    assert quant_config["format"] == quant_format
+    assert sparsity_config["format"] == sparse_format
+
     if os.path.isdir(tmp_path):
         shutil.rmtree(tmp_path)
 
@@ -588,7 +585,7 @@ def _quantization_config_from_string(config_str, q_type):
         quantize_activations=quantize_activations,
         a_bits=a_bits,
         a_type=q_type,
-        a_strategy="channel",
+        a_strategy="tensor",
     )