[dtensor] relax device_mesh argument constraint in local_map #157049
Conversation
This PR relaxes the device_mesh argument constraint in the local_map API. The current restriction is too strict, i.e. all the input arguments must have the same device mesh if they are DTensors. But many times user might want to pass in DTensors to this function that lives on different device mesh, i.e. weight and activation could live in different device mesh. When using the local_map, we are extracting the local tensors from DTensors, and as long as the placements user specified match with the actual DTensor placements, user knows clearly that the inputs are intended to live in different mesh. So this PR removes the same mesh check and update doc to clearly document the behavior. The `device_mesh` argument now serves for a main purpose, allow user to specify the device_mesh for the output DTensor reconstruction
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/157049
Note: Links to docs will display an error until the docs builds have been completed. ❌ 2 New Failures, 1 Cancelled Job, 1 Unrelated FailureAs of commit aa77732 with merge base 81759af ( NEW FAILURES - The following jobs have failed:
CANCELLED JOB - The following job was cancelled. Please retry:
UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
pytorch-bot
bot
added
ciflow/inductor
oncall: distributed
wanchaol
added
release notes: distributed (dtensor)
|
LGTM! Overall this seems to be a hack to get around the cross mesh assertion. We basically leave the responsibility to the user to make sure the shape is correct so that we don't need to write the complicated "cross mesh sharding propagation rule". My concern is that the code can easily fail once we run the model on top of a different world_size or input. |
@zpcore The role of the I think the previous assertion asserts the different input DTensors must have the same mesh is for extra safety in the beginning, but it turns out not necessary. The restriction is that it would limit the expressiveness that the |
|
@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 |
|
@pytorchbot merge -f "ci failure not related" |
Merge startedYour change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
5 participants
This PR relaxes the device_mesh argument constraint in the local_map API. The current restriction is too strict, i.e. all the input arguments must have the same device mesh if they are DTensors. But many times user might want to pass in DTensors to this function that lives on different device mesh, i.e. weight and activation could live in different device mesh.
When using the local_map, we are extracting the local tensors from DTensors, and as long as the placements user specified match with the actual DTensor placements, user knows clearly that the inputs are intended to live in different mesh. So this PR removes the same mesh check and update doc to clearly document the behavior.
The
device_meshargument now serves for a main purpose, allow user to specify the device_mesh for the output DTensor reconstructionFixes #ISSUE_NUMBER
cc @H-Huang @awgu @fegin @fduwjj @wz337 @wconstab @d4l3k