feat(logging): name the offending weight when compression is skipped

changelog.d/85.changed

@@ -0,0 +1 @@
+When a weight is incompatible with the configured quantization block size or palettization granularity, the warning now reports the weight's fully qualified name and shape along with the module name to use in `module_name_configs`, instead of only the compression target.

src/coreai_opt/_utils/insertion/torch_function/handler.py

@@ -9,9 +9,11 @@
 from typing import Any
 
 import torch.nn as nn
+import torch.nn.utils.parametrize as P
 
 from coreai_opt._utils.torch_utils import NamedModule
 from coreai_opt.config import CompressionConfig
+from coreai_opt.config.spec import CompressionSimulatorBase
 
 from .base_supported_ops_registry import BaseSupportedOpsRegistry
 from .modes import (
@@ -21,6 +23,18 @@
 from .types import ModuleCompressionComponents
 
 
+def _record_tensor_fqns(model: nn.Module) -> None:
+    """Record the FQN of each parametrized tensor on the simulator compressing it."""
+    for module_name, module in model.named_modules():
+        if not P.is_parametrized(module):
+            continue
+        for param_name, parametrizations in module.parametrizations.items():
+            fqn = f"{module_name}.{param_name}" if module_name else param_name
+            for simulator in parametrizations:
+                if isinstance(simulator, CompressionSimulatorBase):
+                    simulator.tensor_fqn = fqn
+
+
 class TorchFunctionEagerHandler:
     """
     Prepares the model for compression by inserting weight and activation
@@ -70,6 +84,9 @@ def prepare(self, model: nn.Module, example_inputs: tuple[Any, ...]) -> nn.Modul
         # seen again later in the forward pass.
         register_optimization_mode.register_all_activations()
         register_optimization_mode.register_all_states()
+
+        _record_tensor_fqns(model)
+
         if self._is_weight_only_optimization(self.module_components_dict):
             return model
 

src/coreai_opt/config/spec/compression_simulator.py

@@ -26,6 +26,9 @@ class CompressionSimulatorBase(_ClassRegistryMixin, nn.Module):
     compression simulation is performed during training.
     """
 
+    # FQN of the compressed tensor
+    tensor_fqn: str = "<unknown>"
+
     @abstractmethod
     def forward(self, tensor: torch.Tensor) -> torch.Tensor:
         """

src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py

@@ -166,7 +166,11 @@ def ensure_initialized(self, tensor: torch.Tensor) -> None:
             self._initialize(tensor.detach())
         except (_IncompatibleClusterDimError, _IncompatibleGranularityError) as e:
             logger.warning(
-                f"Tensor incompatible with configured spec: {e}. Skipping palettization."
+                "Tensor '%s' (shape: %s) incompatible with configured spec: %s. "
+                "Skipping palettization.",
+                self.tensor_fqn,
+                tuple(tensor.shape),
+                str(e).rstrip("."),
             )
             self._disabled = True
 

src/coreai_opt/palettization/kmeans/palettizer.py

@@ -64,19 +64,14 @@ class _FakePalettInfo:
     """Metadata about _KMeansFakePalettize module and its associated module / param"""
 
     module: torch.nn.Module
-    module_name: str
     attr_name: str
     idx: int
     fp_module: _KMeansFakePalettize
     weight: torch.Tensor
 
-    @property
-    def layer_name(self) -> str:
-        return f"{self.module_name}.{self.attr_name}" if self.module_name else self.attr_name
-
 
 def _calculate_centroids_for_module(
-    args: tuple[_KMeansFakePalettize, torch.Tensor, str],
+    args: tuple[_KMeansFakePalettize, torch.Tensor],
 ) -> _KMeansFakePalettize:
     """Compute centroids for a single _KMeansFakePalettize module.
 
@@ -87,20 +82,22 @@ def _calculate_centroids_for_module(
     Invokes ``fp_module.forward(weight)`` to mirror the sequential path.
 
     Args:
-        args (tuple[_KMeansFakePalettize, torch.Tensor, str]): Tuple of
-            ``(fp_module, weight, layer_name)``.
+        args (tuple[_KMeansFakePalettize, torch.Tensor]): Tuple of
+            ``(fp_module, weight)``.
 
     Returns:
         _KMeansFakePalettize: The mutated module, ready to be swapped into
             the parent's parametrization slot.
     """
-    fp_module, weight, layer_name = args
+    fp_module, weight = args
 
     try:
         with torch.no_grad():
             fp_module(weight)
     except Exception as e:
-        raise RuntimeError(f"Centroid calculation failed for layer {layer_name!r}") from e
+        raise RuntimeError(
+            f"Centroid calculation failed for weight {fp_module.tensor_fqn!r}"
+        ) from e
 
     return fp_module
 
@@ -505,7 +502,7 @@ def _collect_fake_palett_info(self, *, to_cpu: bool) -> list[_FakePalettInfo]:
         shipped to a spawned worker process.
         """
         fp_info: list[_FakePalettInfo] = []
-        for module_name, module in self._model.named_modules(remove_duplicate=True):
+        for _module_name, module in self._model.named_modules(remove_duplicate=True):
             if not P.is_parametrized(module):
                 continue
             for attr_name, parametrizations in module.parametrizations.items():
@@ -517,7 +514,6 @@ def _collect_fake_palett_info(self, *, to_cpu: bool) -> list[_FakePalettInfo]:
                         fp_info.append(
                             _FakePalettInfo(
                                 module=module,
-                                module_name=module_name,
                                 attr_name=attr_name,
                                 idx=idx,
                                 fp_module=p,
@@ -541,7 +537,7 @@ def _calculate_centroids_sequential(self) -> None:
             return
 
         results = [
-            _calculate_centroids_for_module((info.fp_module, info.weight, info.layer_name))
+            _calculate_centroids_for_module((info.fp_module, info.weight))
             for info in tqdm(fp_info, desc="Palettizing layers (num_workers=1)")
         ]
         self._apply_centroid_results(fp_info, results)
@@ -561,7 +557,7 @@ def _calculate_centroids_parallel(self, num_workers: int) -> None:
 
         # spawn (not fork) so workers don't inherit the parent's CUDA context
         # or other process-global state.
-        pool_args = [(info.fp_module, info.weight, info.layer_name) for info in fp_info]
+        pool_args = [(info.fp_module, info.weight) for info in fp_info]
         ctx = mp.get_context("spawn")
         with ctx.Pool(processes=effective_workers) as pool:
             results = list(
@@ -589,7 +585,7 @@ def _apply_centroid_results(
         """
         for info, new_fp in zip(fp_info, results, strict=True):
             if new_fp.is_disabled():
-                logger.warning("Disabling palettization for layer %r", info.layer_name)
+                logger.warning("Disabling palettization for weight '%s'", new_fp.tensor_fqn)
             info.module.parametrizations[info.attr_name][info.idx] = new_fp
 
     @staticmethod

src/coreai_opt/quantization/_graph/quantizer.py

@@ -46,6 +46,7 @@
 from coreai_opt._utils.version_utils import version_ge
 from coreai_opt.common import ExportBackend
 from coreai_opt.config.compression_config import ModuleConfigDict, _build_module_alias_map
+from coreai_opt.config.spec.compression_simulator import CompressionSimulatorBase
 from coreai_opt.quantization._axis_defaults import (
     apply_weight_axis_defaults_graph as _apply_weight_axis_defaults,
     validate_activation_axes,
@@ -99,6 +100,22 @@
 logger = logging.getLogger(__name__)
 
 
+def _record_tensor_fqns_graph(model: torch.fx.GraphModule) -> None:
+    """Record the FQN of each compressed parameter on the simulator compressing it.
+
+    A simulator reading from anything other than a ``get_attr`` node keeps the
+    default name: an activation, or a weight arriving through a decompression op,
+    has no parameter to name.
+    """
+    simulators = dict(model.named_modules(remove_duplicate=False))
+    for node in model.graph.nodes:
+        if node.op != "call_module" or not node.args:
+            continue
+        simulator = simulators.get(str(node.target))
+        if isinstance(simulator, CompressionSimulatorBase) and node.args[0].op == "get_attr":
+            simulator.tensor_fqn = str(node.args[0].target)
+
+
 class _OpConfigLevel(Enum):
     """
     Enum to specify the op-level config type within a module config.
@@ -1379,6 +1396,8 @@ def _postprocess_prepared_model(model: torch.fx.GraphModule) -> None:
         # granularity is mathematically incorrect. Force per-tensor.
         force_per_tensor_for_channel_altering_ops(model)
 
+        _record_tensor_fqns_graph(model)
+
         # Apply weight axis defaults for per channel and per block quantization
         _apply_weight_axis_defaults(model)
 

src/coreai_opt/quantization/spec/fake_quantize.py

@@ -128,12 +128,14 @@ def enable_observer(self, enabled: bool = True) -> None:
             return
         super().enable_observer(enabled)
 
-    def _warn_and_disable(self, error: _BlockSizeMismatchError) -> None:
-        """Log a warning and permanently disable this module."""
+    def _warn_and_disable(self, error: _BlockSizeMismatchError, shape: torch.Size) -> None:
+        """Log a warning naming the offending tensor and permanently disable this module."""
         logger.warning(
-            "Tensor (target: %s) incompatible with block size "
+            "Tensor '%s' (target: %s, shape: %s) incompatible with block size "
             "configuration: %s. Skipping quantization.",
+            self.tensor_fqn,
             self.quantization_target,
+            tuple(shape),
             error,
         )
         self._disabled.fill_(True)
@@ -156,7 +158,7 @@ def forward(self, tensor: torch.Tensor) -> torch.Tensor:
                 try:
                     scale, zero_point, minval = self.qparams_calculator(tensor)
                 except _BlockSizeMismatchError as e:
-                    self._warn_and_disable(e)
+                    self._warn_and_disable(e, tensor.shape)
                     return tensor
         else:
             # When the observer is not enabled, call the get_qparams

tests/fixtures/quantization.py

@@ -17,6 +17,7 @@
     PerBlockGranularity,
     PerChannelGranularity,
     PerTensorGranularity,
+    QuantizationGranularity,
     QuantizationScheme,
     QuantizationSpec,
 )
@@ -50,23 +51,27 @@ def make_quant_config(
     weight_dtype: torch.dtype | str | None,
     act_dtype: torch.dtype | str | None,
     execution_mode: str,
+    granularity: QuantizationGranularity | None = None,
 ) -> QuantizerConfig:
-    """Build a per-tensor symmetric QuantizerConfig for export tests.
+    """Build a symmetric QuantizerConfig for export tests.
 
     Args:
         weight_dtype (torch.dtype | str | None): Weight dtype, or None to disable.
         act_dtype (torch.dtype | str | None): Activation dtype, or None to disable.
         execution_mode (str): Either "eager" or "graph".
+        granularity (QuantizationGranularity | None): Granularity for both the
+            weight and activation specs. Defaults to per-tensor.
 
     Returns:
-        QuantizerConfig: Config with the requested per-tensor symmetric specs.
+        QuantizerConfig: Config with the requested symmetric specs.
     """
+    granularity = granularity or PerTensorGranularity()
 
     def _spec(dtype: torch.dtype | str) -> QuantizationSpec:
         return QuantizationSpec(
             dtype=dtype,
             qscheme=QuantizationScheme.SYMMETRIC,
-            granularity=PerTensorGranularity(),
+            granularity=granularity,
         )
 
     weight_spec = _spec(weight_dtype) if weight_dtype is not None else None

tests/palettization/test_kmeans_palettizer.py

@@ -296,6 +296,29 @@ def test_disabled_fake_palett_removed_after_prepare(self, caplog):
         assert isinstance(prepared_model[2].parametrizations.weight[0], _FakePalettizeImplBase)
         assert not prepared_model[2].parametrizations.weight[0].is_disabled()
 
+    def test_skip_warning_names_the_offending_weight(self, caplog):
+        """The skip warning identifies the weight by FQN and shape."""
+        # axis 0 is out_features: 12 % 8 != 0.
+        model = nn.Sequential(nn.Linear(8, 12))
+        spec = PalettizationSpec(
+            n_bits=2, granularity=PerGroupedChannelGranularity(axis=0, group_size=8)
+        )
+        config = KMeansPalettizerConfig(
+            global_config=ModuleKMeansPalettizerConfig(op_state_spec={"weight": spec})
+        )
+
+        with caplog.at_level(logging.WARNING):
+            KMeansPalettizer(model, config).prepare((torch.randn(1, 8),))
+
+        skip_messages = [msg for msg in caplog.messages if "Skipping palettization" in msg]
+        assert len(skip_messages) == 1, f"Expected one skip warning, got {skip_messages}"
+        assert skip_messages[0] == (
+            "Tensor '0.weight' (shape: (12, 8)) incompatible with configured spec: Tensor size "
+            "12 along axis 0 is not divisible by group_size 8. For per-grouped-channel "
+            "palettization, the tensor shape along the specified axis must be divisible by "
+            "group_size. Skipping palettization."
+        )
+
     def test_prepared_model_supports_torch_inference(
         self, simple_conv_linear_model, simple_model_input
     ):

tests/quantization/test_eager_quant.py

@@ -38,6 +38,7 @@
     default_weight_quantization_spec,
 )
 from coreai_opt.quantization.spec.fake_quantize import FakeQuantizeImplBase
+from tests.fixtures.quantization import make_quant_config
 
 
 class InnerModule(nn.Module):
@@ -3574,6 +3575,28 @@ def test_non_divisible_block_size_warns_and_skips(self, caplog):
         test_input = torch.randn(1, 768)
         torch.testing.assert_close(prepared_model(test_input), ref_model(test_input))
 
+    def test_skip_warning_names_the_offending_weight(self, caplog):
+        """The skip warning identifies the weight by FQN and shape."""
+        # axis 0 is out_features: 12 % 8 != 0.
+        model = nn.Sequential(nn.Linear(8, 12))
+        config = make_quant_config(
+            weight_dtype="int8",
+            act_dtype=None,
+            execution_mode="eager",
+            granularity=PerBlockGranularity(axis=0, block_size=8),
+        )
+
+        with caplog.at_level(logging.WARNING):
+            Quantizer(model, config).prepare((torch.randn(1, 8),))
+
+        skip_messages = [msg for msg in caplog.messages if "Skipping quantization" in msg]
+        assert len(skip_messages) == 1, f"Expected one skip warning, got {skip_messages}"
+        assert skip_messages[0] == (
+            "Tensor '0.weight' (target: CompressionTargetTensor.WEIGHT, shape: (12, 8)) "
+            "incompatible with block size configuration: Tensor size 12 along axis 0 is not "
+            "divisible by block size 8. Skipping quantization."
+        )
+
     def test_divisible_block_size_not_disabled(self):
         """Linear(768, 1024) with block_size=32 on axis=0: 1024 % 32 == 0."""
         model = nn.Sequential(nn.Linear(768, 1024))

tests/quantization/test_graph_mode_quantizer.py

@@ -55,6 +55,7 @@
     MovingAverageQParamsCalculator,
     StaticQParamsCalculator,
 )
+from tests.fixtures.quantization import make_quant_config
 from tests.models.simple import SimpleModel
 from tests.test_utils.general import get_fake_quant_nodes
 
@@ -1613,6 +1614,28 @@ def test_non_divisible_block_size_warns_and_skips(self, caplog, target, expected
         # The graph must still be runnable after the disabled nodes were removed.
         assert prepared_model(*example_inputs).shape == (1, 1024)
 
+    def test_skip_warning_names_the_offending_weight(self, caplog):
+        """The skip warning identifies the weight by FQN and shape."""
+        # axis 0 is out_features: 12 % 8 != 0.
+        model = torch.nn.Sequential(torch.nn.Linear(8, 12))
+        config = make_quant_config(
+            weight_dtype="int8",
+            act_dtype=None,
+            execution_mode="graph",
+            granularity=PerBlockGranularity(axis=0, block_size=8),
+        )
+
+        with caplog.at_level(logging.WARNING):
+            Quantizer(model, config).prepare((torch.randn(1, 8),))
+
+        skip_messages = [msg for msg in caplog.messages if "Skipping quantization" in msg]
+        assert len(skip_messages) == 1, f"Expected one skip warning, got {skip_messages}"
+        assert skip_messages[0] == (
+            "Tensor '0.weight' (target: CompressionTargetTensor.WEIGHT, shape: (12, 8)) "
+            "incompatible with block size configuration: Tensor size 12 along axis 0 is not "
+            "divisible by block size 8. Skipping quantization."
+        )
+
     @pytest.mark.parametrize(
         "backend",
         [