Fix DCE eliminating random operations by improving is_impure() (#151524) #157981
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/157981
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 11ced11 with merge base 86251ef ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
torch/fx/node.py
Outdated
| } | ||
|
|
||
| if self.target in _random_functions: | ||
| # Only impure if using global RNG (no generator or generator=None) |
There was a problem hiding this comment.
I mean technically they're still impure if generator is used. If I have some calls to rand that use a specific generator and I DCE them, the later generator calls will get modified.
There was a problem hiding this comment.
I thought impure in this context only refers to global side-effects, not local mutations. Am I wrong?
There was a problem hiding this comment.
Edward is right here. it's not about purity of global or local, but whether its a pure function or not.
Especially in the context of fixing this issue (i.e. disrepancy between eager and compiled), we should be treating this as impure
…51524) Random operations with explicit generators can still affect observable behavior when eliminated, causing different generator states between eager and compiled execution modes. All random operations should be preserved by DCE. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
|
@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 |
5 participants
DCE was incorrectly eliminating unused random operations like torch.rand() that have global RNG side effects, causing inconsistent results between eager and compiled execution modes.
Root cause: Python random functions (torch.rand, torch.randn, etc.) don't have the _nondeterministic_seeded attribute, so node.is_impure() returns False, allowing DCE to eliminate them despite advancing global RNG state.
Solution: Enhanced is_impure() in torch/fx/node.py to recognize Python random functions and mark them as impure when they use global RNG, regardless of the impure_random parameter setting. This ensures consistency between eager and compiled execution even when config.fallback_random=False.
Key features:
Tests: Enhanced test_impure_random to verify both FX tracing and AOT compilation codepaths, ensuring random operations are preserved and eager/compiled execution consistency is maintained.
🤖 Generated with Claude Code
Fixes #151524
cc @ezyang @SherlockNoMad @EikanWang @jgong5 @wenzhe-nrv