Fix bug visualizing 1D Tensor using rich #152871
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/152871
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 New FailureAs of commit 1630ad1 with merge base a769114 ( NEW FAILURE - The following job has failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
pytorch-bot
bot
added
the
oncall: distributed
|
@pytorchbot label "release notes: distributed (dtensor)" |
pytorch-bot
bot
added
the
release notes: distributed (dtensor)
wangkuiyi
changed the title
| dtensor_height = shape[0] if len(shape) > 0 else 1 | ||
| dtensor_width = shape[1] if len(shape) > 0 else shape[0] | ||
| dtensor_height = shape[0] | ||
| dtensor_width = shape[1] if len(shape) == 2 else 1 |
There was a problem hiding this comment.
Here I am fixing a bug that I created in the previous PR. When the tensor is 1D, consider it a column vector.
| ) | ||
| for device_index, (shape, offset) in device_shard_shape_and_offsets.items() | ||
| } | ||
|
|
There was a problem hiding this comment.
Here I am fixing another bug. When the tensor is 1D, the shape and offset of each shard is a 1-tuple. As we want to draw them in the 2D screen space, we need to extend each 1-tuple into a 2-tuple. In particular, expand the width of shard to 1, and extend the offset on the x-axis of the screen to be 0.
| """ | ||
| To run the example, use the following command: | ||
| torchrun --standalone --nnodes=1 --nproc-per-node=4 visualize_sharding_example.py | ||
| TERM=xterm-256color torchrun --nproc-per-node=4 visualize_sharding_example.py |
There was a problem hiding this comment.
The environment variable XTERM controls the terminal's coloring capability. Change the default value to xterm-256color to release the power of your terminal app.
|
@pytorchbot merge -f "inductor lint error not related to the PR" |
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 |
|
follow up : #152027 |
Fixes #152848
I didn't fix the bug earlier because the example script didn't exhaustively present all combinations of 1D/2D tensor, 1D/2D mesh, and all possible sharding specs. Therefore, in this PR, I enriched the example script to cover all possible combinations.
cc @H-Huang @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k