diff --git a/src/coreai_opt/quantization/_graph/_annotation_pattern_registry.py b/src/coreai_opt/quantization/_graph/_annotation_pattern_registry.py index aa0c53da..6d422926 100644 --- a/src/coreai_opt/quantization/_graph/_annotation_pattern_registry.py +++ b/src/coreai_opt/quantization/_graph/_annotation_pattern_registry.py @@ -886,9 +886,12 @@ def generate_patterns(cls) -> list[torch.fx.GraphModule]: def generate_qspec_sharing_constraints( cls, node: torch.fx.Node, qspecs: ProvisionalQSpecMap ) -> list[Constraint]: - """Inputs must observe together unless the output is per-channel along the - concat axis, in which case each keeps its own scale and they need only - agree on ``DTYPE``. + """Return the constraints tying the cat's input and output quantizers together. + + They share one quantizer only if every slot (inputs and output) is per-tensor, or + per-channel along a non-concat axis. Otherwise each input keeps its own quantizer, + since the inputs can differ in size along the concat axis, and they only agree on + ``DTYPE``. """ output_slot, input_slots = _shared_op_boundary_slots(node) if not _any_input_slot_populated(input_slots, qspecs): @@ -897,14 +900,15 @@ def generate_qspec_sharing_constraints( all_slots = frozenset([output_slot, *input_slots]) concat_dim = _concat_dim(node) - output_fields = qspecs.get(output_slot, ProvisionalQSpec()).fields - output_granularity = output_fields.get(FieldName.GRANULARITY) - - # A per-tensor granularity has ``axis is None``. - granularity_axis = output_granularity.value.axis if output_granularity is not None else None - per_channel_along_concat_axis = granularity_axis is not None and _same_axis( - granularity_axis, concat_dim, node - ) + per_channel_along_concat_axis = False + for slot in all_slots: + granularity = qspecs.get(slot, ProvisionalQSpec()).fields.get(FieldName.GRANULARITY) + if granularity is None or granularity.value is None: + continue + # A per-tensor granularity has ``axis is None``. + granularity_axis = granularity.value.axis + if granularity_axis is not None and _same_axis(granularity_axis, concat_dim, node): + per_channel_along_concat_axis = True if not per_channel_along_concat_axis: # Sharing an instance subsumes sharing fields, since the merged group # reconciles all of them. diff --git a/tests/quantization/test_annotation_pattern_registry.py b/tests/quantization/test_annotation_pattern_registry.py index 2a6ef05f..f94aaf7e 100644 --- a/tests/quantization/test_annotation_pattern_registry.py +++ b/tests/quantization/test_annotation_pattern_registry.py @@ -402,38 +402,23 @@ def forward(self, x): return torch.cat([x, xt], dim=0) # concat along batch, not channel/spatial -class SimpleCatModel(nn.Module): - """Simple model with concat operation.""" +class CatOfParametersModel(nn.Module): + """Concat of two parameters with different sizes along the concat axis. - def __init__(self): - super().__init__() - self.linear1 = torch.nn.Linear(2, 2, bias=False) - self.linear2 = torch.nn.Linear(2, 2, bias=False) - self.linear3 = torch.nn.Linear(2, 2, bias=False) - - def forward(self, x1, x2): - x1 = self.linear1(x1) - x2 = self.linear2(x2) - x = torch.cat([x1, x2]) - x = self.linear3(x) - return x - - -class SimpleConcatModel(nn.Module): - """Simple model with concat operation.""" + Similar to token concat done in models such as SAM 2's mask decoder: + ``torch.cat([iou_token.weight, mask_tokens.weight], dim=0)``. + """ def __init__(self): super().__init__() - self.linear1 = torch.nn.Linear(2, 2, bias=False) - self.linear2 = torch.nn.Linear(2, 2, bias=False) - self.linear3 = torch.nn.Linear(2, 2, bias=False) + self.token1 = nn.Embedding(1, 4) + self.token2 = nn.Embedding(3, 4) + self.linear = nn.Linear(4, 4) - def forward(self, x1, x2): - x1 = self.linear1(x1) - x2 = self.linear2(x2) - x = torch.concat([x1, x2]) - x = self.linear3(x) - return x + def forward(self, x): + tokens = torch.cat([self.token1.weight, self.token2.weight], dim=0) + out = self.linear(tokens + x) + return out class AddLinearMulModel(nn.Module): @@ -1981,69 +1966,127 @@ def test_shared_observer_preserves_per_block_when_pool_preserves_shape( expected_scale_shape = (2, 4, *spatial[:-1], spatial[-1] // 2) assert tuple(input_fq.qparams_calculator.scale.shape) == expected_scale_shape - def test_cat(self): + @pytest.mark.parametrize("cat_fn", [torch.cat, torch.concat], ids=["cat", "concat"]) + def test_cat_per_tensor_shares_quantizers(self, cat_fn): """ Given a model with - input1 -> relu1 - \ - -> Cat -> Linear - input2 -> relu2 / + input1 -> linear1 + \ + -> cat -> linear3 + input2 -> linear2 / - Check that concat inputs and outputs share the same quantizers. + Check that with per-tensor quantization, the cat's inputs and output share one + quantizer, while linear3's output has its own. + + torch.cat and torch.concat export to different aten ops (aten.cat, aten.concat), + each matched by its own entry in ConcatPattern. """ - model = SimpleCatModel() - inp = (torch.randn(1, 2), torch.randn(1, 2)) + + class CatModel(nn.Module): + def __init__(self): + super().__init__() + self.linear1 = nn.Linear(2, 2, bias=False) + self.linear2 = nn.Linear(2, 2, bias=False) + self.linear3 = nn.Linear(2, 2, bias=False) + + def forward(self, x1, x2): + x = cat_fn([self.linear1(x1), self.linear2(x2)]) + out = self.linear3(x) + return out config = QuantizerConfig( - global_config=ModuleQuantizerConfig(op_state_spec=None, op_input_spec=None), + global_config=ModuleQuantizerConfig( + op_state_spec=None, + op_input_spec=None, + op_output_spec={"*": per_tensor_int8_qspec()}, + ), ) + quantizer = Quantizer(CatModel(), config) + prepared_model = quantizer.prepare(example_inputs=(torch.randn(1, 2), torch.randn(1, 2))) - quantizer = Quantizer(model, config) - prepared_model = quantizer.prepare(example_inputs=inp) + cat_nodes = [ + node + for node in prepared_model.graph.nodes + if node.target in (torch.ops.aten.cat.default, torch.ops.aten.concat.default) + ] + assert len(cat_nodes) == 1 + cat = cat_nodes[0] + cat_input1_quantizer, cat_input2_quantizer = ( + getattr(prepared_model, input_node.target) for input_node in cat.args[0] + ) + (cat_output_quantizer_node,) = cat.users + cat_output_quantizer = getattr(prepared_model, cat_output_quantizer_node.target) + (linear3,) = cat_output_quantizer_node.users + (linear3_output_quantizer_node,) = linear3.users + linear3_output_quantizer = getattr(prepared_model, linear3_output_quantizer_node.target) + + assert cat_input1_quantizer is cat_input2_quantizer + assert cat_input1_quantizer is cat_output_quantizer + assert linear3_output_quantizer is not cat_output_quantizer + # No other quantizers: the one the cat shares, and linear3's output quantizer. quantizers = [ module for module in prepared_model.modules() if isinstance(module, _DefaultFakeQuantizeImpl) ] assert len(quantizers) == 2 - assert prepared_model.activation_post_process_0 is prepared_model.activation_post_process_1 - assert prepared_model.activation_post_process_0 is prepared_model.activation_post_process_2 - # Final linear output quantizer is not associated with the others - assert prepared_model.activation_post_process_0 != prepared_model.activation_post_process_3 + def test_cat_inputs_per_channel_along_concat_axis(self): + """ + Concat inputs that are per-channel along the concat axis keep separate quantizers. - def test_concat(self): + The inputs are parameters named ``weight``, so the default config gives them the + per-channel (axis 0) weight spec, while the output is per-tensor. The inputs differ + in size along axis 0 ([1, 4] and [3, 4]), so one shared per-channel quantizer can't + observe both. """ - Given a model with + model = CatOfParametersModel().eval() + quantizer = Quantizer(model, QuantizerConfig()) + prepared_model = quantizer.prepare(example_inputs=(torch.randn(4, 4),)) - input1 -> relu1 - \ - -> Concat -> Linear - input2 -> relu2 / + cat_nodes = [ + node for node in prepared_model.graph.nodes if node.target == torch.ops.aten.cat.default + ] + assert len(cat_nodes) == 1 + cat = cat_nodes[0] + token1_quantizer, token2_quantizer = ( + getattr(prepared_model, input_node.target) for input_node in cat.args[0] + ) + assert token1_quantizer is not token2_quantizer - Check that concat inputs and outputs share the same quantizers. + def test_cat_inputs_per_channel_along_other_axis(self): """ - model = SimpleConcatModel() - inp = (torch.randn(1, 2), torch.randn(1, 2)) + Concat inputs that are per-channel along a non-concat axis share one quantizer. + The inputs ([1, 4] and [3, 4], concatenated along axis 0) are quantized per-channel + along axis 1, where they have the same size, so one quantizer can observe both. + """ + model = CatOfParametersModel().eval() + weight_spec = QuantizationSpec( + dtype=torch.int8, + qscheme=QuantizationScheme.SYMMETRIC, + granularity=PerChannelGranularity(axis=1), + ) config = QuantizerConfig( - global_config=ModuleQuantizerConfig(op_state_spec=None, op_input_spec=None), + global_config=ModuleQuantizerConfig( + op_state_spec={"weight": weight_spec}, + op_input_spec={"*": per_tensor_int8_qspec()}, + op_output_spec={"*": per_tensor_int8_qspec()}, + ), ) - quantizer = Quantizer(model, config) - prepared_model = quantizer.prepare(example_inputs=inp) - quantizers = [ - module - for module in prepared_model.modules() - if isinstance(module, _DefaultFakeQuantizeImpl) - ] - assert len(quantizers) == 2 - assert prepared_model.activation_post_process_0 is prepared_model.activation_post_process_1 - assert prepared_model.activation_post_process_0 is prepared_model.activation_post_process_2 + prepared_model = quantizer.prepare(example_inputs=(torch.randn(4, 4),)) - # Final linear output quantizer is not associated with the others - assert prepared_model.activation_post_process_0 != prepared_model.activation_post_process_3 + cat_nodes = [ + node for node in prepared_model.graph.nodes if node.target == torch.ops.aten.cat.default + ] + assert len(cat_nodes) == 1 + cat = cat_nodes[0] + token1_quantizer, token2_quantizer = ( + getattr(prepared_model, input_node.target) for input_node in cat.args[0] + ) + assert token1_quantizer is token2_quantizer @pytest.mark.parametrize( "granularity",