[fx] fix split_module with symint #160093
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/160093
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 12d9a77 with merge base 83875cd ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
Test ig? |
Sorry, I didn't understand. |
|
Needs a test! |
|
Sure, I was just checking the patch to see if this negatively affects anything. |
| keep_original_order=True, keep_original_node_name=True) | ||
| return split_gm | ||
|
|
||
| actual = torch.compile(moe, backend=backend)(inp) |
There was a problem hiding this comment.
Couldn't repro the error with torch.fx.symbolic_trace or make_fx.
kshitij12345
changed the title
There was a problem hiding this comment.
Pull Request Overview
This PR fixes a bug in the split_module function where symint dependencies were not being properly tracked across partitions. The fix ensures that when a symbolic integer dependency is detected, the dependent partition properly records the dependency on the partition that defines the symbol.
- Adds missing dependency tracking for symbolic integer nodes across partitions
- Includes a comprehensive test case that reproduces the original issue with GraniteMoe model
Reviewed Changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| torch/fx/passes/split_module.py | Adds dependency tracking for symbolic integer nodes to prevent missing dependencies |
| test/test_fx_experimental.py | Adds test case to verify symint dependency handling in split_module |
| from torch.testing._internal.common_methods_invocations import op_db | ||
| from torch.testing._internal.common_nn import module_tests, get_new_module_tests | ||
| from torch.testing._internal.common_utils import TEST_Z3, run_tests, TestCase | ||
| from torch.testing._internal.common_utils import TEST_Z3, run_tests, TestCase, TEST_WITH_CROSSREF |
There was a problem hiding this comment.
[nitpick] The import statement is mixing single imports with multiple imports. Consider using a consistent import style by importing TEST_WITH_CROSSREF separately or grouping all imports from common_utils on a single line.
Copilot uses AI. Check for mistakes.
|
@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 |
|
@pytorchbot merge |
|
The merge job was canceled or timed out. This most often happen if two merge requests were issued for the same PR, or if merge job was waiting for more than 6 hours for tests to finish. In later case, please do not hesitate to reissue the merge command |
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 |
Fixes #155220 Pull Request resolved: #160093 Approved by: https://github.com/ezyang
Fixes #155220 Pull Request resolved: #160093 Approved by: https://github.com/ezyang
Fixes pytorch#155220 Pull Request resolved: pytorch#160093 Approved by: https://github.com/ezyang
Fixes pytorch#155220 Pull Request resolved: pytorch#160093 Approved by: https://github.com/ezyang
5 participants
Fixes #155220
cc @ezyang @SherlockNoMad @EikanWang @jgong5 @wenzhe-nrv