[DCP] Enable nD device_mesh resharding DTensor in DCP and add associated tests #106230
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/106230
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit e5738e2 with merge base 3fe8417 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
wz337
added
module: distributed_checkpoint
topic: not user facing
wz337
force-pushed
the
test_2d_dtensor_checkpoint
branch
from
fe12fb1 to
73d40f9
Compare
wz337
force-pushed
the
test_2d_dtensor_checkpoint
branch
from
73d40f9 to
100c39d
Compare
wz337
force-pushed
the
test_2d_dtensor_checkpoint
branch
7 times, most recently
from
2db9552 to
dcaf38e
Compare
fix placements change failure add resharding tests
wz337
force-pushed
the
test_2d_dtensor_checkpoint
branch
from
dcaf38e to
e5738e2
Compare
wz337
changed the title
wz337
requested review from
H-Huang,
awgu,
d4l3k,
fduwjj,
fegin,
kiukchung,
kwen2501,
mrshenli,
rohan-varma,
wanchaol and
zhaojuanmao
as code owners
| [Replicate(), Replicate()], | ||
| [Replicate(), Shard(0)], | ||
| [Shard(0), Replicate()], | ||
| [Shard(0), Shard(0)], |
There was a problem hiding this comment.
Maybe consider adding sharding on non-zero dim?
There was a problem hiding this comment.
Thanks for reviewing. Will add more test cases. Sharding on different dim on placements should be no problem. The previous bug is on offset compute is due to when sharding on one dimension twice, dtensor util only gives the local shard offset instead of global offset.
Regardless, will definitely add more test cases for this.
|
@pytorchmergebot 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 |
| for p1 in TWO_D_PLACEMENTS: | ||
| for p2 in TWO_D_PLACEMENTS: |
There was a problem hiding this comment.
nit: you might want to use itertools.product instead?
| global_tensor = torch.arange(16, dtype=torch.float).view(4, 4) | ||
| mesh_shape = (self.world_size,) | ||
| mesh_1d = init_device_mesh(self.device_type, mesh_shape) | ||
| dtensor = distribute_tensor( | ||
| global_tensor, mesh_1d, placements=placements_1d | ||
| ) | ||
| state_dict_to_save = {"dtensor": dtensor} | ||
|
|
||
| dist_cp.save_state_dict( | ||
| state_dict=state_dict_to_save, | ||
| storage_writer=dist_cp.FileSystemWriter(path=CHECKPOINT_DIR), | ||
| planner=dist_cp.DefaultSavePlanner(), | ||
| ) |
There was a problem hiding this comment.
lots of codes here seems to be duplicate. Maybe you want to consolidate them a little bit?
|
LGTM, @fduwjj's comment is legit. Can improve it in a seperate PR. |
5 participants
This PR:
1. Drop assert for 1D DeviceMesh check to allow DTensor with nD DeviceMesh when creating write_item.
2. Add tests for both placement changes and mesh changes for both 1D and 2D scenarios.
cc. @kumpera @wanchaol @fegin