Introduce aggressive merge to CapabilityPartitioner#100195
Introduce aggressive merge to CapabilityPartitioner#100195seanlatias wants to merge 3 commits intopytorch:mainfrom
CapabilityPartitioner#100195Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/100195
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 222e1d2: This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
|
|
@SherlockNoMad please take a look. I tagged you because you are the creater of this pass. Thank you. |
| size = len(merge_candidates_list) | ||
| self_id = merge_candidates_list[0] | ||
| for other_id in merge_candidates_list[1:]: | ||
| maybe_merge_partition(self_id, other_id) |
There was a problem hiding this comment.
I think this is a bug in the original impl. This case should have been handled in line 168. Could you please try fixing the original bug, instead of implementing it as a patch?
There was a problem hiding this comment.
Sure. I can do that.
There was a problem hiding this comment.
@SherlockNoMad I fixed the merging logic. Please take a look and enable the CI. Thanks.
There was a problem hiding this comment.
Thank you for the fix!
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
With the old partitioner, suppose `add` is supported, the following code
```python
def fn(a, b, c, d):
x = a + b # add
y = c + d # add_1
return (x, y)
traced = symbolic_trace(fn)
partitioner = CapabilityBasedPartitioner(traced, supported_ops, allows_single_node_partition=True)
partitions = partitioner.propose_partitions()
```
results in the partitions `[[add], [add_1]]`. However, since these two partitions do not depend on each other, they can be aggressively merged into a single partition `[[add, add_1]]` without causing any issues. This PR introduces a new feature that allows such aggressive merging by introducing an option `aggressive_merge` to the Partitioner class.
Pull Request resolved: pytorch#100195
Approved by: https://github.com/SherlockNoMad
With the old partitioner, suppose
addis supported, the following coderesults in the partitions
[[add], [add_1]]. However, since these two partitions do not depend on each other, they can be aggressively merged into a single partition[[add, add_1]]without causing any issues. This PR introduces a new feature that allows such aggressive merging by introducing an optionaggressive_mergeto the Partitioner class.