diff --git a/src/coreai_opt/quantization/_graph/_annotation_pattern_registry.py b/src/coreai_opt/quantization/_graph/_annotation_pattern_registry.py index 6d42292..66fdc49 100644 --- a/src/coreai_opt/quantization/_graph/_annotation_pattern_registry.py +++ b/src/coreai_opt/quantization/_graph/_annotation_pattern_registry.py @@ -261,7 +261,7 @@ def match_single_pattern( return _annotation_utils.match_pattern_with_subgraph_matcher(model, pattern) -class NAryActPattern(BaseAnnotationPattern[tuple[SourcePartition]]): +class NAryActPattern(BaseAnnotationPattern[tuple[SourcePartition, ...]]): """ Base class for annotation patterns that handle an op of one or more inputs along with an optional activation. @@ -285,7 +285,7 @@ class NAryActPattern(BaseAnnotationPattern[tuple[SourcePartition]]): @classmethod def match_single_pattern( cls, model: torch.fx.GraphModule, pattern: torch.fx.GraphModule | OpsListPattern - ) -> dict[torch.fx.Node, tuple[SourcePartition]]: + ) -> dict[torch.fx.Node, tuple[SourcePartition, ...]]: """ Match nodes in model to the provided pattern and return a dictionary mapping nodes to matches. @@ -300,7 +300,7 @@ def match_single_pattern( return _annotation_utils.match_pattern_with_sequential_partitions(model, pattern) -class SharedObserverModulePattern(BaseAnnotationPattern[tuple[SourcePartition]]): +class SharedObserverModulePattern(BaseAnnotationPattern[tuple[SourcePartition, ...]]): """ Base class for annotation patterns that handle shared observer modules. @@ -353,7 +353,7 @@ def _validate_patterns_length(cls): @classmethod def match_single_pattern( cls, model: torch.fx.GraphModule, pattern: torch.fx.GraphModule | OpsListPattern - ) -> dict[torch.fx.Node, tuple[SourcePartition]]: + ) -> dict[torch.fx.Node, tuple[SourcePartition, ...]]: """ Match nodes in model to the provided pattern and return a dictionary mapping nodes to matches. diff --git a/src/coreai_opt/quantization/_graph/_annotation_utils.py b/src/coreai_opt/quantization/_graph/_annotation_utils.py index 8bbe297..6eaf585 100644 --- a/src/coreai_opt/quantization/_graph/_annotation_utils.py +++ b/src/coreai_opt/quantization/_graph/_annotation_utils.py @@ -19,7 +19,7 @@ SubgraphMatcherWithNameNodeMap, ) from torch.fx.passes.utils.source_matcher_utils import SourcePartition -from torchao.quantization.pt2e import WrapperModule, find_sequential_partitions +from torchao.quantization.pt2e import WrapperModule from torchao.quantization.pt2e.quantizer import ( QuantizationAnnotation, QuantizationSpec as TorchAOQuantizationSpec, @@ -35,6 +35,7 @@ from coreai_opt._utils.fx_utils import ( get_local_state_name, get_module_boundary_nodes, + get_node_type, is_coreai_compressed_state_node, ) from coreai_opt._utils.python_utils import get_fn_arg_names @@ -704,15 +705,6 @@ def _get_call_function_node_from_partition(partition: SourcePartition) -> torch. """ call_function_nodes = [node for node in partition.nodes if node.op == "call_function"] if len(call_function_nodes) != 1: - # torch.export's insert_deferred_runtime_asserts synthesizes one SymInt mul per - # shape-runtime assertion, all sharing one torch_fn tag, so several can collapse - # into a single partition. They carry no tensor value to annotate, so picking any - # one of them is safe here; downstream floating-point filtering no-ops on SymInt. - if call_function_nodes and all( - isinstance(node.meta.get("val"), torch.SymInt) for node in call_function_nodes - ): - return call_function_nodes[0] - module_names = { name for node in call_function_nodes @@ -733,18 +725,89 @@ def _get_call_function_node_from_partition(partition: SourcePartition) -> torch. return call_function_nodes[0] +def _single_op_partition(node: Node, op_type: str) -> SourcePartition: + """Wrap one op node in a ``SourcePartition`` of its own.""" + input_nodes = [arg for arg in node.all_input_nodes if arg.op != "get_attr"] + output_nodes = [node] if node.users else [] + partition = SourcePartition( + nodes=[node], source=op_type, input_nodes=input_nodes, output_nodes=output_nodes + ) + return partition + + +def _find_sequential_partitions( + model: torch.fx.GraphModule, op_types: list[str] +) -> list[tuple[SourcePartition, ...]]: + """Find every chain of ops in ``model`` whose op types are ``op_types``, in order. + + Each chain is returned as a tuple with one ``SourcePartition`` per element of + ``op_types``, each holding exactly one call_function node. + + This is a local replacement for torchao's ``find_sequential_partitions``: it returns + results in the same format, but does not merge unrelated ops into one partition. + + torchao's ``find_sequential_partitions`` gets the ops from torch's + ``get_source_partitions``, which can put separate nodes into a single partition: it + builds partitions based on the first element of ``node.meta["torch_fn"]``, a name that + can be the same for two separate nodes. For example:: + + def forward(self, x1, x2): + s = x1 + x2 # add: torch_fn ("add_1", "method_descriptor.add") + t = torch.add(s, x2) # add_1: torch_fn ("add_1", "builtin_function_or_method.add") + out = torch.relu(t) # relu: only add_1 feeds it + return out + + Both adds get the ``torch_fn`` name ``"add_1"``. Writing each partition as the list of + its nodes, torchao's ``find_sequential_partitions`` returns: + + - for ``["add"]``: one chain, ``([add, add_1],)``: a single partition holding both adds. + - for ``["add", "relu"]``: one chain, ``([add, add_1], [relu])``. + + This function returns: + + - for ``["add"]``: two chains, ``([add],)`` and ``([add_1],)``. + - for ``["add", "relu"]``: one chain, ``([add_1], [relu])``. + + So each add gets its own match, and the inputs of both adds are quantized. + """ + chains = [] + for node in model.graph.nodes: + if node.op == "call_function" and get_node_type(node, warn_on_failure=False) == op_types[0]: + chains.append([node]) + + for op_type in op_types[1:]: + extended_chains = [] + for chain in chains: + for user in chain[-1].users: + if ( + user.op == "call_function" + and get_node_type(user, warn_on_failure=False) == op_type + ): + extended_chains.append([*chain, user]) + chains = extended_chains + + sequences = [] + for chain in chains: + sequence = tuple( + _single_op_partition(node, op_type) + for node, op_type in zip(chain, op_types, strict=True) + ) + sequences.append(sequence) + return sequences + + def match_pattern_with_sequential_partitions( model: torch.fx.GraphModule, pattern: OpsListPattern -) -> dict[torch.fx.Node, tuple[SourcePartition]]: +) -> dict[torch.fx.Node, tuple[SourcePartition, ...]]: """ Helper function for matching model nodes with pattern using - find_sequential_partitions. + ``_find_sequential_partitions``. Returns a dictionary mapping nodes to matches. """ - node_to_match_dict: dict[torch.fx.Node, tuple[SourcePartition]] = {} + node_to_match_dict: dict[torch.fx.Node, tuple[SourcePartition, ...]] = {} - partitions = find_sequential_partitions(model, pattern.pattern) + partitions = _find_sequential_partitions(model, pattern.pattern) # Each sequential partition is a tuple of source partitions, one per pattern element. for sequential_partition in partitions: if _is_branching_partition(sequential_partition): @@ -757,7 +820,7 @@ def match_pattern_with_sequential_partitions( return node_to_match_dict -def _is_branching_partition(sequential_partition: tuple[SourcePartition]) -> bool: +def _is_branching_partition(sequential_partition: tuple[SourcePartition, ...]) -> bool: """ Return True if the partition is a branching partition, False otherwise. diff --git a/tests/quantization/test_annotation_pattern_registry.py b/tests/quantization/test_annotation_pattern_registry.py index f94aaf7..37d1289 100644 --- a/tests/quantization/test_annotation_pattern_registry.py +++ b/tests/quantization/test_annotation_pattern_registry.py @@ -2088,6 +2088,30 @@ def test_cat_inputs_per_channel_along_other_axis(self): ) assert token1_quantizer is token2_quantizer + def test_add_operator_and_add_function(self): + """ + Test that both adds are quantized + """ + + class AddOperatorAndFunctionModel(nn.Module): + def forward(self, x1, x2): + out = x1 + x2 + out = torch.add(out, x2) + return out + + model = AddOperatorAndFunctionModel().eval() + quantizer = Quantizer(model, QuantizerConfig()) + prepared_model = quantizer.prepare(example_inputs=(torch.randn(2, 4), torch.randn(2, 4))) + + patterns = analyze_graph_structure(prepared_model) + check_fake_quantize_placement( + patterns, + { + "add": {"input_fq": [True, True], "output_fq": True}, + "add_1": {"input_fq": [True, True], "output_fq": True}, + }, + ) + @pytest.mark.parametrize( "granularity", [ diff --git a/tests/quantization/test_annotation_utils.py b/tests/quantization/test_annotation_utils.py index 7a630b5..962bf0e 100644 --- a/tests/quantization/test_annotation_utils.py +++ b/tests/quantization/test_annotation_utils.py @@ -10,10 +10,14 @@ import torch.nn as nn import torch.nn.functional as F from torch._ops import OpOverloadPacket -from torchao.quantization.pt2e import WrapperModule +from torchao.quantization.pt2e import ( + WrapperModule, + find_sequential_partitions as torchao_find_sequential_partitions, +) from coreai_opt._utils.torch_utils import export_model from coreai_opt.quantization._graph._annotation_utils import ( + _find_sequential_partitions, _get_aten_graph_module_for_pattern, match_pattern_with_subgraph_matcher, ) @@ -160,3 +164,69 @@ def _relu_pattern(input): model = export_model(nn.ReLU(), (torch.randn(2),), None, False) with pytest.raises(RuntimeError, match='"weight" entry'): match_pattern_with_subgraph_matcher(model, pattern_gm) + + +class TestFindSequentialPartitions: + """Tests for ``_find_sequential_partitions``.""" + + class AddReluMul(nn.Module): + def forward(self, x1, x2): + out = torch.relu(x1 + x2) * x2 + return out + + class AddOperatorAndFunction(nn.Module): + """One add written as ``+`` and one as ``torch.add``: both get ``torch_fn`` "add_1".""" + + def forward(self, x1, x2): + out = x1 + x2 + out = torch.add(out, x2) + return out + + @staticmethod + def _partition_node_names(sequences) -> list[list[list[str]]]: + """Return the node names of each partition, for each sequence.""" + names = [] + for sequence in sequences: + sequence_names = [] + for partition in sequence: + sequence_names.append([node.name for node in partition.nodes]) + names.append(sequence_names) + return names + + @pytest.mark.parametrize( + "op_types", + [["add"], ["add", "relu"], ["relu", "mul"], ["mul"]], + ids=["add", "add_relu", "relu_mul", "mul"], + ) + def test_matches_torchao(self, op_types): + """``_find_sequential_partitions`` mostly returns what torchao's version does. + + They differ only when torch merges unrelated ops into one partition (see + ``test_torch_fn_name_collision``), and in torchao also matching equivalent types, + e.g. "add_" for "add", which the op-list patterns list explicitly instead. + """ + model = export_model( + self.AddReluMul().eval(), (torch.randn(2, 4), torch.randn(2, 4)), None, False + ) + + expected = self._partition_node_names(torchao_find_sequential_partitions(model, op_types)) + actual = self._partition_node_names(_find_sequential_partitions(model, op_types)) + + assert len(actual) == 1 + assert actual == expected + + def test_torch_fn_name_collision(self): + """Each op gets its own partition, even when two ops share a ``torch_fn`` name.""" + model = export_model( + self.AddOperatorAndFunction().eval(), + (torch.randn(2, 4), torch.randn(2, 4)), + None, + False, + ) + + sequences = _find_sequential_partitions(model, ["add"]) + + # torchao's find_sequential_partitions returns [[["add", "add_1"]]] here: one + # sequence whose single partition holds both adds, since torch's + # get_source_partitions groups nodes by torch_fn name. + assert self._partition_node_names(sequences) == [[["add"]], [["add_1"]]] diff --git a/tests/quantization/test_graph_mode_quantizer.py b/tests/quantization/test_graph_mode_quantizer.py index 7d761a5..49966a8 100644 --- a/tests/quantization/test_graph_mode_quantizer.py +++ b/tests/quantization/test_graph_mode_quantizer.py @@ -555,9 +555,9 @@ def test_prepare_with_symint_mul_partition_collision(self): torch.export's ``insert_deferred_runtime_asserts`` synthesizes one SymInt ``mul`` per shape-runtime assertion, all sharing one ``torch_fn`` tag, so - they collapse into a single ``SourcePartition``. The annotator picks the - first node and the downstream ``_is_fx_node_floating_point`` filter no-ops - on SymInt inputs. + they collapse into a single ``SourcePartition``. The annotator splits it + into one partition per node, and the downstream ``_is_fx_node_floating_point`` + filter no-ops on the SymInt inputs. """ H, W, B, embed_dim = 4, 4, 1, 8