support as_python_constant on PlacementClassVariable #124398
Conversation
[ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/124398
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (1 Unrelated Failure)As of commit 945073e with merge base e16f1ee ( BROKEN TRUNK - The following job failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
github-actions
bot
requested review from
SherlockNoMad,
albanD,
antoniojkim,
ezyang and
miladm
[ghstack-poisoned]
| return False | ||
|
|
||
| from torch.distributed._tensor.placement_types import Placement | ||
|
|
There was a problem hiding this comment.
@wanchaol Can you remind me why we needed this variable?
There was a problem hiding this comment.
This boils down to the "custom fx type" that we've been discussing. So fundamentally we need to inline Placement and DeviceMesh as a constant in the closure function (i.e. from_local) and put that closure function to the fx graph as a call_function node.
But if we don't have this PlacementClassVariable, dynamo would trace the DTensor's metadata construction as a UserDefinedClass/Object (i.e. Shard(1)). UDTs are not constant variable (as it's hard to tell whether a UDT is a ConstantVariable unless user explicitly tell dynamo the objects are ConstantVariable). So this PlacementClassVariable allow us to turn the Shard(1) as a PlacementVariable (which is a constant variable), and then the sharding metadata can be inlined as a closure
|
@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 |
Merge failedReason: 1 mandatory check(s) failed. The first few are: Dig deeper by viewing the failures on hud |
Fixes an error for torchtitan + internal cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang [ghstack-poisoned]
Fixes an error for torchtitan + internal cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang [ghstack-poisoned]
Fixes #125287 Fixes #124090, context on the issue Pull Request resolved: #124399 Approved by: https://github.com/soulitzer ghstack dependencies: #124398
…or is noncontiguous (#124400) Fixes #124397 Pull Request resolved: #124400 Approved by: https://github.com/ezyang, https://github.com/yoyoyocmu ghstack dependencies: #124398, #124399
…spatch__ (#123347)" (#125288) Re-land of #123347. The original PR broke internal because of a circular import due to importing dynamo in the DTensor code. The new version uses `torch._dynamo_disable` to work around This reverts commit 9d88339. Pull Request resolved: #125288 Approved by: https://github.com/ezyang, https://github.com/yanboliang, https://github.com/yoyoyocmu, https://github.com/anijain2305, https://github.com/fegin ghstack dependencies: #124398, #124399, #124400
Fixes an error for torchtitan + internal Pull Request resolved: #124398 Approved by: https://github.com/ezyang, https://github.com/wanchaol, https://github.com/yoyoyocmu
Fixes #125287 Fixes #124090, context on the issue Pull Request resolved: #124399 Approved by: https://github.com/soulitzer ghstack dependencies: #124398
…or is noncontiguous (pytorch#124400) Fixes pytorch#124397 Pull Request resolved: pytorch#124400 Approved by: https://github.com/ezyang, https://github.com/yoyoyocmu ghstack dependencies: pytorch#124398, pytorch#124399
…spatch__ (pytorch#123347)" (pytorch#125288) Re-land of pytorch#123347. The original PR broke internal because of a circular import due to importing dynamo in the DTensor code. The new version uses `torch._dynamo_disable` to work around This reverts commit 9d88339. Pull Request resolved: pytorch#125288 Approved by: https://github.com/ezyang, https://github.com/yanboliang, https://github.com/yoyoyocmu, https://github.com/anijain2305, https://github.com/fegin ghstack dependencies: pytorch#124398, pytorch#124399, pytorch#124400
5 participants
Fixes an error for torchtitan + internal
Stack from ghstack (oldest at bottom):
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @chauhang