Activation Ordering Strategies (#146)

Author: kylesayrs

src/compressed_tensors/compressors/base.py

@@ -108,13 +108,15 @@ def compress(
                 prefix = name[: -(len(weight_suffix))]
                 scale = model_state.get(merge_names(prefix, "weight_scale"), None)
                 zp = model_state.get(merge_names(prefix, "weight_zero_point"), None)
+                g_idx = model_state.get(merge_names(prefix, "weight_g_idx"), None)
                 if scale is not None:
                     # weight is quantized, compress it
                     quant_args = names_to_scheme[prefix]
                     compressed_data = self.compress_weight(
                         weight=value,
                         scale=scale,
                         zero_point=zp,
+                        g_idx=g_idx,
                         quantization_args=quant_args,
                         device="cpu",
                     )

src/compressed_tensors/quantization/lifecycle/initialize.py

@@ -21,6 +21,7 @@
     wrap_module_forward_quantized,
 )
 from compressed_tensors.quantization.quant_args import (
+    ActivationOrdering,
     QuantizationArgs,
     QuantizationStrategy,
 )
@@ -179,8 +180,8 @@ def _initialize_scale_zero_point_observer(
         )
         module.register_parameter(f"{base_name}_zero_point", init_zero_point)
 
-    # initialize with empty for actorder, to be populated by GPTQ or state_dict
-    if quantization_args.actorder:
+    # only grouped activation ordering has g_idx
+    if quantization_args.actorder == ActivationOrdering.GROUP:
         g_idx_shape = (weight_shape[1],)
         g_idx_dtype = torch.int
         init_g_idx = Parameter(

src/compressed_tensors/quantization/quant_args.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from enum import Enum
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, Union
 
 import torch
 from pydantic import BaseModel, Field, field_validator, model_validator
@@ -25,6 +25,7 @@
     "QuantizationStrategy",
     "QuantizationArgs",
     "round_to_quantized_type",
+    "ActivationOrdering",
 ]
 
 FP8_DTYPE = torch.float8_e4m3fn
@@ -51,6 +52,19 @@ class QuantizationStrategy(str, Enum):
     TOKEN = "token"
 
 
+class ActivationOrdering(str, Enum):
+    """
+    Enum storing strategies for activation ordering
+
+    Group: reorder groups and weight\n
+    Weight: only reorder weight, not groups. Slightly lower latency and
+    accuracy compared to group actorder\n
+    """
+
+    GROUP = "group"
+    WEIGHT = "weight"
+
+
 class QuantizationArgs(BaseModel, use_enum_values=True):
     """
     User facing arguments used to define a quantization config for weights or
@@ -69,17 +83,17 @@ class QuantizationArgs(BaseModel, use_enum_values=True):
         quantization. Note that enabling dynamic quantization will change the default
         observer to a memoryless one
     :param actorder: whether to apply group quantization in decreasing order of
-        activation. Defaults to False for arbitrary ordering
+        activation. Defaults to None for arbitrary ordering
     """
 
     num_bits: int = 8
-    type: QuantizationType = QuantizationType.INT.value
+    type: QuantizationType = QuantizationType.INT
     symmetric: bool = True
     group_size: Optional[int] = None
     strategy: Optional[QuantizationStrategy] = None
     block_structure: Optional[str] = None
     dynamic: bool = False
-    actorder: bool = False
+    actorder: Optional[ActivationOrdering] = None
     observer: str = Field(
         default="minmax",
         description=(
@@ -108,8 +122,15 @@ def get_observer(self):
 
         return Observer.load_from_registry(self.observer, quantization_args=self)
 
+    @field_validator("type", mode="before")
+    def validate_type(cls, value) -> QuantizationType:
+        if isinstance(value, str):
+            return QuantizationType(value.lower())
+
+        return value
+
     @field_validator("group_size", mode="before")
-    def validate_group(cls, value) -> int:
+    def validate_group(cls, value) -> Union[int, None]:
         if value is None:
             return value
 
@@ -121,18 +142,29 @@ def validate_group(cls, value) -> int:
 
         return value
 
-    @model_validator(mode="before")
-    def validate_strategy(values) -> Dict[str, Any]:
-        model_fields = QuantizationArgs.model_fields
-        strategy = values.get("strategy", model_fields["strategy"].default)
-        group_size = values.get("group_size", model_fields["group_size"].default)
-        actorder = values.get("actorder", model_fields["actorder"].default)
+    @field_validator("strategy", mode="before")
+    def validate_strategy(cls, value) -> Union[QuantizationStrategy, None]:
+        if isinstance(value, str):
+            return QuantizationStrategy(value.lower())
 
-        if strategy is not None:
-            strategy = QuantizationStrategy(strategy.lower())
+        return value
 
-        else:
-            # use group_size to determinine strategy if not given explicity
+    @field_validator("actorder", mode="before")
+    def validate_actorder(cls, value) -> Optional[ActivationOrdering]:
+        if isinstance(value, str):
+            return ActivationOrdering(value.lower())
+
+        return value
+
+    @model_validator(mode="after")
+    def validate_model_after(model: "QuantizationArgs") -> Dict[str, Any]:
+        # extract user-passed values from dictionary
+        strategy = model.strategy
+        group_size = model.group_size
+        actorder = model.actorder
+
+        # infer strategy
+        if strategy is None:
             if group_size is None:
                 strategy = QuantizationStrategy.TENSOR
             elif group_size > 0:
@@ -145,21 +177,24 @@ def validate_strategy(values) -> Dict[str, Any]:
                     "strategy='group' and group_size = -1 for 'channel'"
                 )
 
+        # validate strategy and group
         if strategy == QuantizationStrategy.GROUP:
             if group_size is None or group_size <= 0:
                 raise ValueError(
                     f"strategy {strategy} requires group_size to be "
                     "set to a positive value"
                 )
 
-        if actorder and strategy != QuantizationStrategy.GROUP:
+        # validate activation ordering and strategy
+        if actorder is not None and strategy != QuantizationStrategy.GROUP:
             raise ValueError(
-                "Group quantization must be specified in order to apply "
+                "Must use group quantization strategy in order to apply "
                 "activation ordering"
             )
 
-        values["strategy"] = strategy
-        return values
+        # write back modified values
+        model.strategy = strategy
+        return model
 
     def pytorch_dtype(self) -> torch.dtype:
         if self.type == QuantizationType.FLOAT:

tests/test_compressors/test_pack_quant.py

@@ -33,33 +33,32 @@
     apply_quantization_status,
 )
 from compressed_tensors.quantization.lifecycle.forward import fake_quantize
+from compressed_tensors.quantization.quant_args import ActivationOrdering
 from safetensors.torch import save_file
 from torch.nn.modules import Linear, Sequential
 
 
-def get_dummy_quant_config(num_bits=4, strategy=None, group_size=None):
+def get_dummy_quant_config(num_bits=4, strategy=None, group_size=None, actorder=None):
     config_groups = {
         "group_1": QuantizationScheme(
             targets=["Linear"],
             weights=QuantizationArgs(
                 num_bits=num_bits,
                 strategy=strategy,
                 group_size=group_size,
+                actorder=actorder,
             ),
         ),
     }
-    ignore = ["lm_head"]
-    quant_config = QuantizationConfig(
-        config_groups=config_groups,
-        ignore=ignore,
-    )
-
-    return quant_config
+    return QuantizationConfig(config_groups=config_groups)
 
 
 def make_dummy_g_idx(columns: int, group_size: int) -> torch.Tensor:
     perm = torch.randperm(columns)
-    return torch.tensor([index // group_size for index in range(columns)])[perm]
+    return torch.nn.Parameter(
+        (torch.arange(columns, dtype=torch.int) // group_size)[perm],
+        requires_grad=False,
+    )
 
 
 @pytest.mark.parametrize(
@@ -199,29 +198,34 @@ def test_reload_match(tmp_path, num_bits):
 
 
 @pytest.mark.parametrize(
-    "apply_gptq",
-    [True, False],
+    "actorder",
+    [
+        ActivationOrdering.GROUP,
+        ActivationOrdering.WEIGHT,
+        None,
+    ],
 )
-def test_actorder_reload_match(apply_gptq, tmp_path):
-    model = Sequential(
-        OrderedDict(
-            [
-                ("dummy", Linear(512, 1024, bias=None)),
-            ]
-        )
-    )
+def test_actorder_reload_match(actorder, tmp_path):
+    model = Sequential(OrderedDict([("dummy", Linear(512, 1024, bias=None))]))
     group_size = 128
-    quant_config = get_dummy_quant_config(strategy="group", group_size=group_size)
+    quant_config = get_dummy_quant_config(
+        strategy="group", group_size=group_size, actorder=actorder
+    )
     apply_quantization_config(model, quant_config)
-    apply_quantization_status(model, QuantizationStatus.CALIBRATION)
-
-    if apply_gptq:
-        model.dummy.weight_g_idx = make_dummy_g_idx(512, group_size)
 
+    # run calibration
+    apply_quantization_status(model, QuantizationStatus.CALIBRATION)
     for _ in range(16):
         inputs = torch.rand((512, 512))
         _ = model(inputs)
+    apply_quantization_status(model, QuantizationStatus.FROZEN)
+
+    # apply gptq
+    if actorder == ActivationOrdering.GROUP:
+        init_g_idx = make_dummy_g_idx(512, group_size)
+        model.dummy.register_parameter("weight_g_idx", init_g_idx)
 
+    # compress
     compressor = PackedQuantizationCompressor(config=quant_config)
     quantized_modules_to_args = {
         "dummy": quant_config.config_groups["group_1"].weights,
@@ -230,6 +234,8 @@ def test_actorder_reload_match(apply_gptq, tmp_path):
         model.state_dict(), names_to_scheme=quantized_modules_to_args
     )
     save_file(compressed_state_dict, tmp_path / "model.safetensors")
+
+    # decompress
     reconstructed_dense_gen = compressor.decompress(
         tmp_path, names_to_scheme=quantized_modules_to_args
     )
@@ -241,6 +247,7 @@ def test_actorder_reload_match(apply_gptq, tmp_path):
         model.dummy.weight,
         scale=model.dummy.weight_scale,
         zero_point=model.dummy.weight_zero_point,
+        g_idx=getattr(model.dummy, "weight_g_idx", None),
         args=quantized_modules_to_args["dummy"],
     )
     assert torch.equal(fake_quant_dummy, reconstructed_dense["dummy.weight"])

tests/test_quantization/test_quant_args.py

@@ -14,6 +14,7 @@
 
 import pytest
 from compressed_tensors.quantization import (
+    ActivationOrdering,
     QuantizationArgs,
     QuantizationStrategy,
     QuantizationType,
@@ -39,6 +40,9 @@ def test_group():
     assert group.strategy == QuantizationStrategy.GROUP
     assert group.group_size == kwargs["group_size"]
 
+    with pytest.raises(ValueError):
+        QuantizationArgs(strategy=QuantizationStrategy.GROUP, group_size=-1)
+
 
 def test_block():
     kwargs = {"strategy": "block", "block_structure": "2x4"}
@@ -56,25 +60,40 @@ def test_infer_strategy():
     assert args.strategy == QuantizationStrategy.CHANNEL
 
 
+def test_enums():
+    assert QuantizationArgs(
+        type=QuantizationType.INT,
+        strategy=QuantizationStrategy.GROUP,
+        actorder=ActivationOrdering.WEIGHT,
+        group_size=1,
+    ) == QuantizationArgs(type="InT", strategy="GROUP", actorder="weight", group_size=1)
+
+
 def test_actorder():
-    args = QuantizationArgs(group_size=128, actorder=True)
+    # test group inference with actorder
+    args = QuantizationArgs(group_size=128, actorder=ActivationOrdering.GROUP)
     assert args.strategy == QuantizationStrategy.GROUP
-    assert args.actorder
 
+    # test invalid pairings
     with pytest.raises(ValueError):
-        args = QuantizationArgs(group_size=None, actorder=True)
-
+        QuantizationArgs(group_size=None, actorder="weight")
     with pytest.raises(ValueError):
-        args = QuantizationArgs(group_size=-1, actorder=True)
-
+        QuantizationArgs(group_size=-1, actorder="weight")
     with pytest.raises(ValueError):
-        args = QuantizationArgs(strategy="tensor", actorder=True)
+        QuantizationArgs(strategy="tensor", actorder="weight")
+
+    # test boolean defaulting
+    assert (
+        QuantizationArgs(group_size=1, actorder="weight").actorder
+        == ActivationOrdering.WEIGHT
+    )
+    assert QuantizationArgs(group_size=1, actorder=None).actorder is None
 
 
 def test_invalid():
     with pytest.raises(ValidationError):
-        _ = QuantizationArgs(type="invalid")
+        QuantizationArgs(type="invalid")
     with pytest.raises(ValidationError):
-        _ = QuantizationArgs(strategy="invalid")
+        QuantizationArgs(strategy="invalid")
     with pytest.raises(ValidationError):
-        _ = QuantizationArgs(strategy=QuantizationStrategy.GROUP)
+        QuantizationArgs(strategy=QuantizationStrategy.GROUP)