diff --git a/changelog.d/85.changed b/changelog.d/85.changed new file mode 100644 index 00000000..11db1aea --- /dev/null +++ b/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. diff --git a/src/coreai_opt/_utils/insertion/torch_function/handler.py b/src/coreai_opt/_utils/insertion/torch_function/handler.py index a5f1682c..9225a532 100644 --- a/src/coreai_opt/_utils/insertion/torch_function/handler.py +++ b/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 diff --git a/src/coreai_opt/config/spec/compression_simulator.py b/src/coreai_opt/config/spec/compression_simulator.py index d3324e23..d47465b7 100644 --- a/src/coreai_opt/config/spec/compression_simulator.py +++ b/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 = "" + @abstractmethod def forward(self, tensor: torch.Tensor) -> torch.Tensor: """ diff --git a/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py b/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py index 566c140a..18404b11 100644 --- a/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py +++ b/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 diff --git a/src/coreai_opt/palettization/kmeans/palettizer.py b/src/coreai_opt/palettization/kmeans/palettizer.py index b85d7c64..e0aa27d8 100644 --- a/src/coreai_opt/palettization/kmeans/palettizer.py +++ b/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 diff --git a/src/coreai_opt/quantization/_graph/quantizer.py b/src/coreai_opt/quantization/_graph/quantizer.py index 339968b5..44181ab5 100644 --- a/src/coreai_opt/quantization/_graph/quantizer.py +++ b/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) diff --git a/src/coreai_opt/quantization/spec/fake_quantize.py b/src/coreai_opt/quantization/spec/fake_quantize.py index 9ff72aa9..31b4423e 100644 --- a/src/coreai_opt/quantization/spec/fake_quantize.py +++ b/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 diff --git a/tests/fixtures/quantization.py b/tests/fixtures/quantization.py index 84ea8f7b..511ad399 100644 --- a/tests/fixtures/quantization.py +++ b/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 diff --git a/tests/palettization/test_kmeans_palettizer.py b/tests/palettization/test_kmeans_palettizer.py index 5379ee8e..a152998e 100644 --- a/tests/palettization/test_kmeans_palettizer.py +++ b/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 ): diff --git a/tests/quantization/test_eager_quant.py b/tests/quantization/test_eager_quant.py index a4bccd5d..5bf65eec 100644 --- a/tests/quantization/test_eager_quant.py +++ b/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)) diff --git a/tests/quantization/test_graph_mode_quantizer.py b/tests/quantization/test_graph_mode_quantizer.py index 71cda955..6f3c8a14 100644 --- a/tests/quantization/test_graph_mode_quantizer.py +++ b/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", [