[DONOTREVIEW][DTenosr][Test] DTensor 2D sharding

[wz337]

world_size = 6
mesh_shape(3, 2)
placements = (Shard(0), Shard(0))

print result:

self.rank=0, dtensor.to_local().shape=torch.Size([2]), dtensor.to_local()=tensor([0, 1], device='cuda:0')
self.rank=1, dtensor.to_local().shape=torch.Size([1]), dtensor.to_local()=tensor([2], device='cuda:1')
self.rank=2, dtensor.to_local().shape=torch.Size([2]), dtensor.to_local()=tensor([3, 4], device='cuda:2')
self.rank=3, dtensor.to_local().shape=torch.Size([1]), dtensor.to_local()=tensor([5], device='cuda:3')
self.rank=4, dtensor.to_local().shape=torch.Size([1]), dtensor.to_local()=tensor([6], device='cuda:4')
self.rank=5, dtensor.to_local().shape=torch.Size([1]), dtensor.to_local()=tensor([7], device='cuda:5')

Right now, the sharding result is:

rank0: [0, 1], rank1: [2]
rank2: [3, 4], rank3: [5]
rank4: [6], rank5: [7] 

This result is expected as DTensor don't handle padding right now.

What is expected with padding should be:

rank0: [0, 1], rank1: [4, 5]
rank2: [2, 3], rank3: [6, 7]
rank4: [], rank5: [] 

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @XilunWu @wanchaol @fduwjj @tianyu-l @wconstab @yf225 @chauhang @d4l3k

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/124339

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❌ 3 New Failures, 1 Unrelated Failure

As of commit c859a2d with merge base dd3cea3 ():

NEW FAILURES - The following jobs have failed:

• inductor / cuda12.1-py3.10-gcc9-sm86 / test (inductor, 1, 1, linux.g5.4xlarge.nvidia.gpu) (gh)
  test/distributed/tensor/parallel/test_fsdp_2d_parallel.py::Test2DReshardingLoad::test_load_state_dict
• Lint / lintrunner-noclang / linux-job (gh)
  >>> Lint for test/distributed/tensor/parallel/test_fsdp_2d_parallel.py:
• pull / linux-jammy-py3.8-gcc11 / test (distributed, 1, 2, linux.2xlarge) (gh)
  distributed/tensor/parallel/test_fsdp_2d_parallel.py::Test2DReshardingLoad::test_load_state_dict

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / linux-focal-cuda12.1-py3.10-gcc9 / test (default, 5, 5, linux.4xlarge.nvidia.gpu) (gh)
  test_foreach.py::TestForeachCUDA::test_binary_op_list_error_cases__foreach_add_cuda_complex128

This comment was automatically generated by Dr. CI and updates every 15 minutes.