[DTensor] add explicit mode (ExplicitRedistributionContext)#166593
[DTensor] add explicit mode (ExplicitRedistributionContext)#166593wconstab wants to merge 3 commits intogh/wconstab/448/basefrom
Conversation
[ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/166593
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 9eab1bf with merge base 397d9fe ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
pytorch-bot
bot
added
the
oncall: distributed
|
This is at least better than using An op is hard to target from user script, because unlike parameters, they don't have FQNs. So a user would either have to use a graph based approach, where the ops become nodes (which you would have a handle to). But this approach requires the user to use graph passes to parallelize the program. Or the user can pass through the coat of nn.Module, and directly interact with the ops eagerly, like you did. |
cc H-Huang awgu wanchaol fegin fduwjj wz337 d4l3k pragupta msaroufim dcci [ghstack-poisoned]
cc H-Huang awgu wanchaol fegin fduwjj wz337 d4l3k pragupta msaroufim dcci [ghstack-poisoned]
wconstab
changed the title
|
@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 |
…torch#167370) Also support nesting, enable/disable, and make the class use a thread-local for storage so independent threads do not confuse each other. Pull Request resolved: pytorch#167370 Approved by: https://github.com/ezyang ghstack dependencies: pytorch#166593
…166593) usage: ``` dx = distribute_tensor(x, device_mesh, [Shard(0)]) dA = distribute_tensor(A, device_mesh, [Shard(0)]) with ExplicitRedistributionContext(): with self.assertRaisesRegex(RuntimeError, "Implicit redistribution"): # Shard(0) @ Shard(0) requires a redistribution torch.matmul(dx, dA) ``` Pull Request resolved: pytorch#166593 Approved by: https://github.com/ezyang
…torch#167370) Also support nesting, enable/disable, and make the class use a thread-local for storage so independent threads do not confuse each other. Pull Request resolved: pytorch#167370 Approved by: https://github.com/ezyang ghstack dependencies: pytorch#166593
4 participants
Stack from ghstack (oldest at bottom):
usage:
cc @H-Huang @awgu @wanchaol @fegin @fduwjj @wz337 @d4l3k @pragupta @msaroufim @dcci