Pass Dependencies When Proposing Partitions VIA New Group-Based Partitioner #12072
Conversation
leafs1
requested review from
JacobSzwejbka,
digantdesai,
jackzhxng,
larryliu0820,
manuelcandales,
mcr229,
mergennachin and
swolchok
as code owners
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/12072
Note: Links to docs will display an error until the docs builds have been completed. ❌ 2 Cancelled JobsAs of commit 115cdfa with merge base c778063 ( CANCELLED JOBS - The following jobs were cancelled. Please retry:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
facebook-github-bot
added
the
CLA Signed
leafs1
force-pushed
the
partitioner
branch
2 times, most recently
from
8397b92 to
afefd1d
Compare
|
@pytorchbot label "release notes: none" |
pytorch-bot
bot
added
the
release notes: none
mergennachin
added
release notes: exir
Still not fully grasped the motivation, can you explain this a bit more perhaps with an example? And why the existing greedy partitioner merging logic from the CapabilityBasedParitioner would split these nodes? Thanks. |
digantdesai
left a comment
digantdesai
left a comment
There was a problem hiding this comment.
Looks good, left some comments. Mainly around rationale and testing. I assume we will start using this in the XNNPACKPartitioner, if yes, that might help with testing.
| allowed_single_node_partition_ops: Optional[Sequence[str]] = None, | ||
| node_groups: List[List[Node]] = None, | ||
| ) -> None: | ||
| super().__init__( |
|
|
||
| return False | ||
|
|
||
| def find_nodes_by_names( |
There was a problem hiding this comment.
| def find_nodes_by_names( | |
| def _find_nodes_by_names( |
| return result | ||
|
|
||
| def create_model(self, model): | ||
| return model() |
There was a problem hiding this comment.
| return model() | |
| return model().eval() |
| [set(node_group) for node_group in node_groups] if node_groups else None | ||
| ) | ||
| self.node_to_group = collections.defaultdict(int) | ||
| self.all_nodes_in_groups = set() |
There was a problem hiding this comment.
what if a node is in two groups? Should we check against that?
| if self.node_groups: | ||
| for i, group in enumerate(self.node_groups): | ||
| # Create a partition for each group | ||
| partition_id = next(new_partition_id) |
There was a problem hiding this comment.
What is the rationale for not allowing two node_groups in the same partition? Or can they get merged later, still reading...
| ) | ||
| self.check_partition( | ||
| partitions, {"add", "linear_2", "fake_quantize_per_tensor_affine_1"} | ||
| ) |
There was a problem hiding this comment.
I don't believe these tests are sufficient to cover the rather huge propose_partitions method. I propose, in addition to these, can we a lot more tests with node_groups = None and compare the partitioned output with CapabilityBasedPartitioner?
leafs1
force-pushed
the
partitioner
branch
2 times, most recently
from
bfb6bda to
5901086
Compare
Here's the initial github issue with an example #11447. The idea is that the current greedy partitioner when dealing with a quantized graph can partition the dq/q nodes in different partitions (example partition 1 has q and partition 2 has the op and the dq). As a result, the semantics of the quantized op is lost in the graph. |
mcr229
left a comment
mcr229
left a comment
There was a problem hiding this comment.
@leafs1 thanks for all these extra tests michael. I do believe it is well tested. Just as a follow up let's test for faulty cases (enforced groups might not be valid (because of dependencies)). And let's also put up a PR wiring this up to xnnpack so we can see how this fairs.
Not sure if this is added as well. But let's put some tests up where compatibility and group based partitioner differ in their outputted partition. Specifically in a quantization use case where a quantized node is moved out of its necessary dependency.
5 participants
Summary
The existing capability-based partitioner in PyTorch's FX module lacks the ability to specify node dependencies that must be partitioned together. This can lead to incorrect partitioning of dynamically quantized linear patterns, resulting in the loss of quantization semantics. For example, in a graph with shared QDQ (Quantize-Dequantize) chains, the partitioner may incorrectly separate nodes that should remain together, leading to incorrect execution semantics.
This PR addresses this issue by adding a new group-based partitioner that allows users to specify groups of nodes that must be partitioned together.
Test plan
I've created test cases that replicate existing QDQ tests, as well as additional graphs with different node dependencies. These tests aim to verify that the new partitioner correctly groups nodes together based on the specified dependencies.