[Utils] Skip internal modules when matching (#404)

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

Author: kylesayrs

src/compressed_tensors/quantization/lifecycle/forward.py

@@ -112,17 +112,21 @@ def dequantize(
             if scale.shape[1] == 1:
                 args = QuantizationArgs(strategy=QuantizationStrategy.CHANNEL)
             # Scale height matches input or is 1 -> group quantization across columns
-            # 
+            #
             # Example 1: scale.shape[0] == 1
             # x_q: (4, 8), scale: (1, 4) -> 2 columns per group
             #
-            # Example 2: scale.shape[0] == x_q.shape[0] 
+            # Example 2: scale.shape[0] == x_q.shape[0]
             # x_q: (4, 8), scale: (4, 4) -> 2 elements per group (per row)
             elif (scale.shape[0] == 1) or (scale.shape[0] == x_q.shape[0]):
                 group_size = int(x_q.shape[1] / scale.shape[1])
-                args = QuantizationArgs(strategy=QuantizationStrategy.GROUP, group_size=group_size)
+                args = QuantizationArgs(
+                    strategy=QuantizationStrategy.GROUP, group_size=group_size
+                )
             else:
-                args = QuantizationArgs(strategy=QuantizationStrategy.BLOCK, block_structure=scale.shape)
+                args = QuantizationArgs(
+                    strategy=QuantizationStrategy.BLOCK, block_structure=scale.shape
+                )
         else:
             raise ValueError(
                 f"Could not infer a quantization strategy from scale with {scale.ndim} "

src/compressed_tensors/quantization/lifecycle/initialize.py

@@ -185,27 +185,29 @@ def _initialize_scale_zero_point(
         elif quantization_args.strategy == QuantizationStrategy.BLOCK:
             # For block quantization, scale shape should match number of blocks - only for weights
             if quantization_args.block_structure is None:
-                raise ValueError("Block quantization requires block_structure to be specified")
+                raise ValueError(
+                    "Block quantization requires block_structure to be specified"
+                )
             block_height, block_width = quantization_args.block_structure
             rows, cols = weight_shape[-2], weight_shape[-1]
             num_rows_blocks = math.ceil(rows / block_height)
             num_cols_blocks = math.ceil(cols / block_width)
-            
+
             # Warn if dimensions don't divide evenly
             if rows % block_height != 0 or cols % block_width != 0:
                 warnings.warn(
                     f"Block quantization: tensor shape {weight_shape} does not divide evenly "
                     f"by block structure {quantization_args.block_structure}. "
                     f"Some blocks will be incomplete which may affect quantization quality.",
-                    UserWarning
+                    UserWarning,
                 )
-            
+
             expected_shape = (num_rows_blocks, num_cols_blocks)
     elif quantization_args.strategy == QuantizationStrategy.BLOCK:
         warnings.warn(
             f"BLOCK quantization not supported for {base_name} activations. "
             f"Falling back to tensor-level quantization.",
-            UserWarning
+            UserWarning,
         )
         expected_shape = 1
 

src/compressed_tensors/quantization/quant_scheme.py

@@ -64,8 +64,9 @@ def validate_model_after(model: "QuantizationScheme") -> "QuantizationScheme":
                 raise ValueError("Cannot apply actorder to output activations")
 
         if (
-            inputs and weights
-            and weights.strategy == QuantizationStrategy.GROUP 
+            inputs
+            and weights
+            and weights.strategy == QuantizationStrategy.GROUP
             and inputs.strategy == QuantizationStrategy.GROUP
             and weights.group_size != inputs.group_size
         ):
@@ -75,7 +76,7 @@ def validate_model_after(model: "QuantizationScheme") -> "QuantizationScheme":
                 "may complicate fused kernel implementations. Consider using "
                 "TENSOR_GROUP strategy for both or matching group sizes.",
                 UserWarning,
-                stacklevel=2
+                stacklevel=2,
             )
 
         return model

src/compressed_tensors/transform/factory/hadamard.py

@@ -12,9 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import math
 from typing import Optional, Union
 
-import math
 import torch
 from compressed_tensors.transform import TransformArgs, TransformScheme
 from compressed_tensors.transform.factory.base import TransformBase, TransformFactory
@@ -103,7 +103,8 @@ def forward(self, value: Tensor) -> Tensor:
 
         if self.args.inverse:
             weight = weight.T
- 
-        return apply_transform_weight(
-            weight, value, self.args.location, self.module_type
-        ) / self._scale
+
+        return (
+            apply_transform_weight(weight, value, self.args.location, self.module_type)
+            / self._scale
+        )

src/compressed_tensors/utils/match.py

@@ -18,6 +18,7 @@
 from typing import Iterable, Tuple
 
 import torch
+from compressed_tensors.utils.internal import InternalModule
 
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
@@ -28,8 +29,6 @@
     "match_named_parameters",
     "match_modules_set",
     "is_match",
-    "match_name",
-    "match_class",
 ]
 
 
@@ -83,13 +82,16 @@ def match_named_parameters(
     """
     unmatched_targets = set(targets)
     for module_name, module in model.named_modules():
+        if isinstance(module, InternalModule):
+            continue
+
         for param_name, param in module.named_parameters(recurse=False):
             param_fqn = f"{module_name}.{param_name}"
             for target in targets:
-                if match_name(param_fqn, target):
+                if _match_name(param_fqn, target):
                     unmatched_targets -= {target}
 
-                    if not any(match_name(param_fqn, ign) for ign in ignore):
+                    if not any(_match_name(param_fqn, ign) for ign in ignore):
                         yield param_fqn, module, param
 
     if warn_on_fail:
@@ -165,11 +167,14 @@ def match_modules_set(
 def is_match(name: str, module: torch.nn.Module, target: str) -> bool:
     """
     Returns true if either module name or module parent classes match against target
+    and the module is not an internal module
     """
-    return match_name(name, target) or match_class(module, target)
+    return not isinstance(module, InternalModule) and (
+        _match_name(name, target) or _match_class(module, target)
+    )
 
 
-def match_name(name: str, target: str) -> bool:
+def _match_name(name: str, target: str) -> bool:
     """
     Returns true if target string begins with "re:" and
     regex matches or if target string exactly matches name
@@ -180,7 +185,7 @@ def match_name(name: str, target: str) -> bool:
         return target == name
 
 
-def match_class(module: torch.nn.Module, target: str) -> bool:
+def _match_class(module: torch.nn.Module, target: str) -> bool:
     """
     Returns true if any torch parent class names match the target string exactly
     """

tests/test_utils/test_match.py

@@ -19,14 +19,14 @@
 from accelerate import init_empty_weights
 
 # Assuming the module is named "module_matching" - adjust import as needed
-from compressed_tensors.utils.match import (
+from compressed_tensors.utils import (
+    InternalModule,
     is_match,
-    match_class,
     match_modules_set,
-    match_name,
     match_named_modules,
     match_named_parameters,
 )
+from compressed_tensors.utils.match import _match_class, _match_name
 
 
 class DummyModel(nn.Module):
@@ -66,14 +66,14 @@ def __init__(self):
 
 
 class TestMatchName:
-    """Test cases for match_name function"""
+    """Test cases for _match_name function"""
 
     def test_exact_match(self):
         """Test exact string matching"""
-        assert match_name("layer1", "layer1") == True
-        assert match_name("layer1", "layer2") == False
+        assert _match_name("layer1", "layer1") == True
+        assert _match_name("layer1", "layer2") == False
         assert (
-            match_name(
+            _match_name(
                 "transformer.layers.0.self_attn.q_proj",
                 "transformer.layers.0.self_attn.q_proj",
             )
@@ -82,14 +82,14 @@ def test_exact_match(self):
 
     def test_regex_match(self):
         """Test regex matching with "re:" prefix"""
-        assert match_name("layer1", "re:layer.*") == True
-        assert match_name("layer1", "re:^layer1$") == True
-        assert match_name("layer1", "re:layer2") == False
+        assert _match_name("layer1", "re:layer.*") == True
+        assert _match_name("layer1", "re:^layer1$") == True
+        assert _match_name("layer1", "re:layer2") == False
         assert (
-            match_name("transformer.layers.0.self_attn.q_proj", "re:.*q_proj") == True
+            _match_name("transformer.layers.0.self_attn.q_proj", "re:.*q_proj") == True
         )
         assert (
-            match_name(
+            _match_name(
                 "transformer.layers.0.self_attn.q_proj",
                 "re:transformer\\.layers\\.\\d+\\.self_attn\\..*_proj$",
             )
@@ -98,49 +98,49 @@ def test_regex_match(self):
 
     def test_empty_strings(self):
         """Test edge cases with empty strings"""
-        assert match_name("", "") == True
-        assert match_name("layer1", "") == False
-        assert match_name("", "layer1") == False
+        assert _match_name("", "") == True
+        assert _match_name("layer1", "") == False
+        assert _match_name("", "layer1") == False
 
     def test_regex_special_characters(self):
         """Test regex with special characters"""
-        assert match_name("layer.1", "re:layer\\.1") == True
-        assert match_name("layer.1", "re:layer.1") == True  # . matches any char
-        assert match_name("layer_1", "re:layer_1") == True
+        assert _match_name("layer.1", "re:layer\\.1") == True
+        assert _match_name("layer.1", "re:layer.1") == True  # . matches any char
+        assert _match_name("layer_1", "re:layer_1") == True
 
 
 class TestMatchClass:
-    """Test cases for match_class function"""
+    """Test cases for _match_class function"""
 
     def test_direct_class_match(self):
         """Test matching direct class names"""
         linear = nn.Linear(10, 20)
-        assert match_class(linear, "Linear") == True
-        assert match_class(linear, "Conv2d") == False
+        assert _match_class(linear, "Linear") == True
+        assert _match_class(linear, "Conv2d") == False
 
         norm = nn.LayerNorm(10)
-        assert match_class(norm, "LayerNorm") == True
-        assert match_class(norm, "BatchNorm1d") == False
+        assert _match_class(norm, "LayerNorm") == True
+        assert _match_class(norm, "BatchNorm1d") == False
 
     def test_parent_class_match(self):
         """Test matching parent class names"""
         linear = nn.Linear(10, 20)
-        assert match_class(linear, "Module") == True
+        assert _match_class(linear, "Module") == True
 
         conv = nn.Conv2d(3, 16, 3)
-        assert match_class(conv, "Module") == True
-        assert match_class(conv, "_ConvNd") == True
+        assert _match_class(conv, "Module") == True
+        assert _match_class(conv, "_ConvNd") == True
 
     def test_non_torch_module(self):
         """Test with non-torch modules"""
         regular_object = object()
-        assert match_class(regular_object, "object") == False  # not a torch.nn.Module
+        assert _match_class(regular_object, "object") == False  # not a torch.nn.Module
 
     def test_custom_module(self):
         """Test with custom module classes"""
         model = DummyModel()
-        assert match_class(model, "DummyModel") == True
-        assert match_class(model, "Module") == True
+        assert _match_class(model, "DummyModel") == True
+        assert _match_class(model, "Module") == True
 
 
 class TestIsMatch:
@@ -171,6 +171,15 @@ def test_regex_in_name_match(self):
         assert is_match("layer1", linear, "re:layer.*") == True
         assert is_match("layer1", linear, "re:conv.*") == False
 
+    def test_internal_module_match(self):
+        """Test not matching internal modules"""
+
+        class InternalLinear(InternalModule, nn.Linear):
+            pass
+
+        linear = InternalLinear(10, 20)
+        assert is_match("layer1", linear, "re:layer.*") == False
+
 
 class TestMatchNamedModules:
     """Test cases for match_named_modules function"""
@@ -236,6 +245,16 @@ def test_warn_on_fail(self, mock_logger):
         assert "Could not match" in warning_msg
         assert "nonexistent_module" in warning_msg
 
+    def test_internal_match(self):
+        """Test not matching internal modules"""
+
+        class InternalLinear(InternalModule, nn.Linear):
+            pass
+
+        linear = InternalLinear(10, 20)
+        matches = list(match_named_modules(linear, ["re:.*"]))
+        assert len(matches) == 0
+
 
 class TestMatchNamedParameters:
     """Test cases for match_named_parameters function"""
@@ -298,6 +317,16 @@ def test_warn_on_fail_parameters(self, mock_logger):
         assert "Could not match" in warning_msg
         assert "nonexistent.param" in warning_msg
 
+    def test_internal_match(self):
+        """Test not matching internal modules"""
+
+        class InternalLinear(InternalModule, nn.Linear):
+            pass
+
+        linear = InternalLinear(10, 20)
+        matches = list(match_named_parameters(linear, ["re:.*"]))
+        assert len(matches) == 0
+
 
 class TestMatchModulesSet:
     """Test cases for match_modules_set function"""
@@ -377,6 +406,16 @@ def test_module_set_with_ignore(self):
         # Should have 2 sets (layers 1 and 2, but not 0)
         assert len(matches) == 2
 
+    def test_internal_match(self):
+        """Test not matching internal modules"""
+
+        class InternalLinear(InternalModule, nn.Linear):
+            pass
+
+        linear = InternalLinear(10, 20)
+        matches = list(match_modules_set(linear, ["re:.*"]))
+        assert len(matches) == 0
+
 
 class TestIntegration:
     """Integration tests combining multiple functions"""