From 8eba42fff3e4cddee01d01d4b9adeb8907ebb856 Mon Sep 17 00:00:00 2001 From: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com> Date: Tue, 7 Jul 2026 16:46:08 -0700 Subject: [PATCH 1/3] refactor: use dataclass for weight axis defaults table Signed-off-by: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com> --- src/coreai_opt/quantization/_axis_defaults.py | 90 +++++++++---------- tests/quantization/test_axis_defaults.py | 7 +- 2 files changed, 45 insertions(+), 52 deletions(-) diff --git a/src/coreai_opt/quantization/_axis_defaults.py b/src/coreai_opt/quantization/_axis_defaults.py index 6db37fe..99ff143 100644 --- a/src/coreai_opt/quantization/_axis_defaults.py +++ b/src/coreai_opt/quantization/_axis_defaults.py @@ -13,6 +13,7 @@ import logging from collections import defaultdict +from dataclasses import dataclass import torch import torch.nn as nn @@ -35,43 +36,45 @@ _WeightFQMap = defaultdict[FakeQuantizeImplBase, list[_ConsumerInfo]] -# Per-channel defaults: quantize along the OUTPUT channel axis. -# Conv weights are [out_ch, in_ch, ...]; ConvTranspose weights are [in_ch, out_ch, ...]. -# Embedding weights are [num_embeddings, embedding_dim]. -_PER_CHANNEL_WEIGHT_AXIS_DEFAULTS: dict[type[nn.Module], int] = { - nn.Conv1d: 0, - nn.Conv2d: 0, - nn.Conv3d: 0, - nn.ConvTranspose1d: 1, - nn.ConvTranspose2d: 1, - nn.ConvTranspose3d: 1, - nn.Linear: 0, - nn.Embedding: 0, -} +@dataclass(frozen=True) +class _WeightAxisSpec: + """Default axis indices and aten op mapping for a single ``nn.Module`` type. + + Attributes: + per_channel_axis (int): Output-channel axis for per-channel granularity. + per_block_axis (int): Input-channel (reduction) axis for per-block granularity. + aten_ops (tuple[torch._ops.OpOverload, ...]): Graph-mode aten ops that + correspond to this module type. + """ + + per_channel_axis: int + per_block_axis: int + aten_ops: tuple[torch._ops.OpOverload, ...] -# Per-block defaults: quantize along the INPUT channel axis (reduction dimension) + def default_axis_for(self, granularity: QuantizationGranularity) -> int | None: + """Return the default axis for the given granularity, or ``None`` if not applicable.""" + if isinstance(granularity, PerChannelGranularity): + return self.per_channel_axis + if isinstance(granularity, PerBlockGranularity): + return self.per_block_axis + return None + + +# Conv weights are [out_ch, in_ch, ...]; ConvTranspose weights are [in_ch, out_ch, ...]. # Embedding weights are [num_embeddings, embedding_dim]. -_PER_BLOCK_WEIGHT_AXIS_DEFAULTS: dict[type[nn.Module], int] = { - nn.Conv1d: 1, - nn.Conv2d: 1, - nn.Conv3d: 1, - nn.ConvTranspose1d: 0, - nn.ConvTranspose2d: 0, - nn.ConvTranspose3d: 0, - nn.Linear: 1, - nn.Embedding: 1, +_WEIGHT_AXIS_SPECS: dict[type[nn.Module], _WeightAxisSpec] = { + nn.Conv1d: _WeightAxisSpec(0, 1, (torch.ops.aten.conv1d.default,)), + nn.Conv2d: _WeightAxisSpec(0, 1, (torch.ops.aten.conv2d.default,)), + nn.Conv3d: _WeightAxisSpec(0, 1, (torch.ops.aten.conv3d.default,)), + nn.ConvTranspose1d: _WeightAxisSpec(1, 0, (torch.ops.aten.conv_transpose1d.default,)), + nn.ConvTranspose2d: _WeightAxisSpec(1, 0, (torch.ops.aten.conv_transpose2d.input,)), + nn.ConvTranspose3d: _WeightAxisSpec(1, 0, (torch.ops.aten.conv_transpose3d.input,)), + nn.Linear: _WeightAxisSpec(0, 1, (torch.ops.aten.linear.default,)), + nn.Embedding: _WeightAxisSpec(0, 1, (torch.ops.aten.embedding.default,)), } -# Maps aten OpOverload -> nn.Module type for graph-mode op to module type resolution. _ATEN_OP_TO_MODULE_TYPE: dict[torch._ops.OpOverload, type[nn.Module]] = { - torch.ops.aten.conv1d.default: nn.Conv1d, - torch.ops.aten.conv2d.default: nn.Conv2d, - torch.ops.aten.conv3d.default: nn.Conv3d, - torch.ops.aten.conv_transpose1d.default: nn.ConvTranspose1d, - torch.ops.aten.conv_transpose2d.input: nn.ConvTranspose2d, - torch.ops.aten.conv_transpose3d.input: nn.ConvTranspose3d, - torch.ops.aten.linear.default: nn.Linear, - torch.ops.aten.embedding.default: nn.Embedding, + op: mod for mod, spec in _WEIGHT_AXIS_SPECS.items() for op in spec.aten_ops } @@ -220,12 +223,11 @@ def _collect_weight_fq_entries_eager(model: nn.Module) -> _WeightFQMap: def _apply_defaults(fq_map: _WeightFQMap) -> None: """Apply default weight axes and raise on unresolved or conflicting entries. - For each FQ in the map, selects the per-channel or per-block defaults - table based on the granularity type. Resolves each consumer's identifier - to a module type (aten ``OpOverload`` entries from graph-mode are mapped via - ``_ATEN_OP_TO_MODULE_TYPE``), then looks up the default axis. When a - single FQ has multiple consumers (shared weight), all consumers must - agree on the same default axis. + For each FQ in the map, resolves each consumer's identifier to a module + type (aten ``OpOverload`` entries from graph-mode are mapped via + ``_ATEN_OP_TO_MODULE_TYPE``), then looks up the default axis from + ``_WEIGHT_AXIS_SPECS``. When a single FQ has multiple consumers (shared + weight), all consumers must agree on the same default axis. Args: fq_map (_WeightFQMap): Map from fake-quantize instance to its @@ -243,14 +245,6 @@ def _apply_defaults(fq_map: _WeightFQMap) -> None: if not _granularity_needs_axis_default(fq.granularity): continue - # get appropriate defaults table - if isinstance(fq.granularity, PerChannelGranularity): - axis_defaults = _PER_CHANNEL_WEIGHT_AXIS_DEFAULTS - elif isinstance(fq.granularity, PerBlockGranularity): - axis_defaults = _PER_BLOCK_WEIGHT_AXIS_DEFAULTS - else: - continue - # Resolve each consumer to a module type and look up its default axis resolved_axes: set[int] = set() all_consumer_names: list[str] = [] @@ -261,8 +255,8 @@ def _apply_defaults(fq_map: _WeightFQMap) -> None: else: module_type = module_type_or_op - # pick a default based on this consumer - default_axis = axis_defaults.get(module_type) if module_type is not None else None + spec = _WEIGHT_AXIS_SPECS.get(module_type) if module_type is not None else None + default_axis = spec.default_axis_for(fq.granularity) if spec is not None else None if default_axis is not None: resolved_axes.add(default_axis) diff --git a/tests/quantization/test_axis_defaults.py b/tests/quantization/test_axis_defaults.py index 950fa1d..a370d4e 100644 --- a/tests/quantization/test_axis_defaults.py +++ b/tests/quantization/test_axis_defaults.py @@ -19,8 +19,7 @@ QuantizerConfig, ) from coreai_opt.quantization._axis_defaults import ( - _PER_BLOCK_WEIGHT_AXIS_DEFAULTS, - _PER_CHANNEL_WEIGHT_AXIS_DEFAULTS, + _WEIGHT_AXIS_SPECS, _apply_defaults, _WeightFQMap, ) @@ -184,9 +183,9 @@ def test_axis_none_resolved( prepared = Quantizer(model, config).prepare((make_input(),)) if isinstance(granularity, PerChannelGranularity): - expected_axis = _PER_CHANNEL_WEIGHT_AXIS_DEFAULTS[module_type] + expected_axis = _WEIGHT_AXIS_SPECS[module_type].per_channel_axis else: - expected_axis = _PER_BLOCK_WEIGHT_AXIS_DEFAULTS[module_type] + expected_axis = _WEIGHT_AXIS_SPECS[module_type].per_block_axis weight_fqs = _get_weight_fqs(prepared) assert len(weight_fqs) == 1 From 1743357f3c64c75660b62050b4524a76b22fed49 Mon Sep 17 00:00:00 2001 From: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com> Date: Thu, 9 Jul 2026 18:14:44 -0700 Subject: [PATCH 2/3] address comments --- src/coreai_opt/_utils/torch_utils.py | 13 ++++++++ src/coreai_opt/quantization/_axis_defaults.py | 30 ++++++++----------- tests/quantization/test_axis_defaults.py | 20 +++++++++++++ 3 files changed, 45 insertions(+), 18 deletions(-) diff --git a/src/coreai_opt/_utils/torch_utils.py b/src/coreai_opt/_utils/torch_utils.py index 915891c..465d9ea 100644 --- a/src/coreai_opt/_utils/torch_utils.py +++ b/src/coreai_opt/_utils/torch_utils.py @@ -33,6 +33,19 @@ E8M0_EXPONENT_BIAS: Final[int] = 127 F32_MIN_NORMAL: Final[float] = 2**-126 # ~1.175e-38 +# Maps a graph-mode aten ``OpOverload`` to the ``nn.Module`` type it corresponds +# to. +ATEN_OP_TO_MODULE_TYPE: dict[torch._ops.OpOverload, type[torch.nn.Module]] = { + torch.ops.aten.conv1d.default: torch.nn.Conv1d, + torch.ops.aten.conv2d.default: torch.nn.Conv2d, + torch.ops.aten.conv3d.default: torch.nn.Conv3d, + torch.ops.aten.conv_transpose1d.default: torch.nn.ConvTranspose1d, + torch.ops.aten.conv_transpose2d.input: torch.nn.ConvTranspose2d, + torch.ops.aten.conv_transpose3d.input: torch.nn.ConvTranspose3d, + torch.ops.aten.linear.default: torch.nn.Linear, + torch.ops.aten.embedding.default: torch.nn.Embedding, +} + class NamedModule(NamedTuple): """NamedTuple for holding name and module info""" diff --git a/src/coreai_opt/quantization/_axis_defaults.py b/src/coreai_opt/quantization/_axis_defaults.py index 99ff143..e973082 100644 --- a/src/coreai_opt/quantization/_axis_defaults.py +++ b/src/coreai_opt/quantization/_axis_defaults.py @@ -20,6 +20,7 @@ import torch.nn.utils.parametrize as P from torch.fx import GraphModule +from coreai_opt._utils.torch_utils import ATEN_OP_TO_MODULE_TYPE from coreai_opt.config.spec import CompressionTargetTensor from coreai_opt.quantization.spec.fake_quantize import FakeQuantizeImplBase from coreai_opt.quantization.spec.granularity import ( @@ -38,18 +39,15 @@ @dataclass(frozen=True) class _WeightAxisSpec: - """Default axis indices and aten op mapping for a single ``nn.Module`` type. + """Default axis indices for a single ``nn.Module`` type. Attributes: per_channel_axis (int): Output-channel axis for per-channel granularity. per_block_axis (int): Input-channel (reduction) axis for per-block granularity. - aten_ops (tuple[torch._ops.OpOverload, ...]): Graph-mode aten ops that - correspond to this module type. """ per_channel_axis: int per_block_axis: int - aten_ops: tuple[torch._ops.OpOverload, ...] def default_axis_for(self, granularity: QuantizationGranularity) -> int | None: """Return the default axis for the given granularity, or ``None`` if not applicable.""" @@ -63,18 +61,14 @@ def default_axis_for(self, granularity: QuantizationGranularity) -> int | None: # Conv weights are [out_ch, in_ch, ...]; ConvTranspose weights are [in_ch, out_ch, ...]. # Embedding weights are [num_embeddings, embedding_dim]. _WEIGHT_AXIS_SPECS: dict[type[nn.Module], _WeightAxisSpec] = { - nn.Conv1d: _WeightAxisSpec(0, 1, (torch.ops.aten.conv1d.default,)), - nn.Conv2d: _WeightAxisSpec(0, 1, (torch.ops.aten.conv2d.default,)), - nn.Conv3d: _WeightAxisSpec(0, 1, (torch.ops.aten.conv3d.default,)), - nn.ConvTranspose1d: _WeightAxisSpec(1, 0, (torch.ops.aten.conv_transpose1d.default,)), - nn.ConvTranspose2d: _WeightAxisSpec(1, 0, (torch.ops.aten.conv_transpose2d.input,)), - nn.ConvTranspose3d: _WeightAxisSpec(1, 0, (torch.ops.aten.conv_transpose3d.input,)), - nn.Linear: _WeightAxisSpec(0, 1, (torch.ops.aten.linear.default,)), - nn.Embedding: _WeightAxisSpec(0, 1, (torch.ops.aten.embedding.default,)), -} - -_ATEN_OP_TO_MODULE_TYPE: dict[torch._ops.OpOverload, type[nn.Module]] = { - op: mod for mod, spec in _WEIGHT_AXIS_SPECS.items() for op in spec.aten_ops + nn.Conv1d: _WeightAxisSpec(0, 1), + nn.Conv2d: _WeightAxisSpec(0, 1), + nn.Conv3d: _WeightAxisSpec(0, 1), + nn.ConvTranspose1d: _WeightAxisSpec(1, 0), + nn.ConvTranspose2d: _WeightAxisSpec(1, 0), + nn.ConvTranspose3d: _WeightAxisSpec(1, 0), + nn.Linear: _WeightAxisSpec(0, 1), + nn.Embedding: _WeightAxisSpec(0, 1), } @@ -225,7 +219,7 @@ def _apply_defaults(fq_map: _WeightFQMap) -> None: For each FQ in the map, resolves each consumer's identifier to a module type (aten ``OpOverload`` entries from graph-mode are mapped via - ``_ATEN_OP_TO_MODULE_TYPE``), then looks up the default axis from + ``ATEN_OP_TO_MODULE_TYPE``), then looks up the default axis from ``_WEIGHT_AXIS_SPECS``. When a single FQ has multiple consumers (shared weight), all consumers must agree on the same default axis. @@ -251,7 +245,7 @@ def _apply_defaults(fq_map: _WeightFQMap) -> None: for module_type_or_op, name in consumers: # graph-mode consumers are aten ops, so map them to module types first if isinstance(module_type_or_op, torch._ops.OpOverload): - module_type = _ATEN_OP_TO_MODULE_TYPE.get(module_type_or_op) + module_type = ATEN_OP_TO_MODULE_TYPE.get(module_type_or_op) else: module_type = module_type_or_op diff --git a/tests/quantization/test_axis_defaults.py b/tests/quantization/test_axis_defaults.py index a370d4e..e09fd76 100644 --- a/tests/quantization/test_axis_defaults.py +++ b/tests/quantization/test_axis_defaults.py @@ -21,6 +21,7 @@ from coreai_opt.quantization._axis_defaults import ( _WEIGHT_AXIS_SPECS, _apply_defaults, + _WeightAxisSpec, _WeightFQMap, ) from coreai_opt.quantization.spec import default_activation_quantization_spec @@ -292,6 +293,25 @@ def _activation_only_config( ) +class TestWeightAxisSpec: + """Verify _WeightAxisSpec.default_axis_for maps granularity types to axes.""" + + def test_per_channel_returns_per_channel_axis(self): + """PerChannelGranularity resolves to the spec's per_channel_axis.""" + spec = _WeightAxisSpec(per_channel_axis=0, per_block_axis=1) + assert spec.default_axis_for(PerChannelGranularity(axis=None)) == 0 + + def test_per_block_returns_per_block_axis(self): + """PerBlockGranularity resolves to the spec's per_block_axis.""" + spec = _WeightAxisSpec(per_channel_axis=0, per_block_axis=1) + assert spec.default_axis_for(PerBlockGranularity(axis=None, block_size=2)) == 1 + + def test_other_granularity_returns_none(self): + """A granularity that is neither per-channel nor per-block returns None.""" + spec = _WeightAxisSpec(per_channel_axis=0, per_block_axis=1) + assert spec.default_axis_for(PerTensorGranularity()) is None + + class TestActivationAxisValidation: """Verify that activation FQs with unresolved axis=None are caught at prepare time.""" From d8afd658ac51c55c7c90d2c63048893dacb76e96 Mon Sep 17 00:00:00 2001 From: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com> Date: Thu, 9 Jul 2026 18:18:32 -0700 Subject: [PATCH 3/3] import privately Signed-off-by: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com> --- src/coreai_opt/quantization/_axis_defaults.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/coreai_opt/quantization/_axis_defaults.py b/src/coreai_opt/quantization/_axis_defaults.py index e973082..cabc0e0 100644 --- a/src/coreai_opt/quantization/_axis_defaults.py +++ b/src/coreai_opt/quantization/_axis_defaults.py @@ -20,7 +20,9 @@ import torch.nn.utils.parametrize as P from torch.fx import GraphModule -from coreai_opt._utils.torch_utils import ATEN_OP_TO_MODULE_TYPE +from coreai_opt._utils.torch_utils import ( + ATEN_OP_TO_MODULE_TYPE as _ATEN_OP_TO_MODULE_TYPE, +) from coreai_opt.config.spec import CompressionTargetTensor from coreai_opt.quantization.spec.fake_quantize import FakeQuantizeImplBase from coreai_opt.quantization.spec.granularity import ( @@ -219,7 +221,7 @@ def _apply_defaults(fq_map: _WeightFQMap) -> None: For each FQ in the map, resolves each consumer's identifier to a module type (aten ``OpOverload`` entries from graph-mode are mapped via - ``ATEN_OP_TO_MODULE_TYPE``), then looks up the default axis from + ``_ATEN_OP_TO_MODULE_TYPE``), then looks up the default axis from ``_WEIGHT_AXIS_SPECS``. When a single FQ has multiple consumers (shared weight), all consumers must agree on the same default axis. @@ -245,7 +247,7 @@ def _apply_defaults(fq_map: _WeightFQMap) -> None: for module_type_or_op, name in consumers: # graph-mode consumers are aten ops, so map them to module types first if isinstance(module_type_or_op, torch._ops.OpOverload): - module_type = ATEN_OP_TO_MODULE_TYPE.get(module_type_or_op) + module_type = _ATEN_OP_TO_MODULE_TYPE.get(module_type_or_op) else: module_type = module_type_or_op