trunk/87ceb3fbef7fbc34d98350587fd2584a615c6dfc: [DTensor] Support _StridedShard to Shard through all-to-all (#170915)

(AI generated commit description)

[DTensor] Support _StridedShard to Shard through all-to-all

Summary

This PR adds support for redistributing tensors from _StridedShard placement to Shard placement using the all-to-all collective operation.

The key challenge is that _StridedShard produces non-contiguous (interleaved) shards, so converting to a regular Shard placement requires:

1. Properly computing padding for both the source strided dimension and target dimension
2. Reordering elements after the all-to-all to restore contiguous layout

Example: Converting _StridedShard(0, split_factor=2) to Shard(1)

Consider the following setup:

• Mesh shape: (4,) — 4 ranks on a single mesh dimension
• Original tensor shape: (9, 4)
• Source placement: (_StridedShard(0, split_factor=2),)
• Target placement: (Shard(1),)

Step 1: Understand the _StridedShard distribution

With _StridedShard(0, split_factor=2), the tensor is conceptually split in two levels on dimension 0:

1. First level: Split into split_factor=2 pieces → chunks of size ⌈9/2⌉ = 5, giving pieces [0:5] and [5:9]
2. Second level: Each piece is split into num_chunks=4 pieces (mesh size)

The shards are then interleaved so each rank gets one slice from each first-level piece:

Original tensor (9x4):            Strided sharding on dim 0:
┌─────────────────────┐
│ row 0               │  ─┐
│ row 1               │   ├─ First piece [0:5], split into 4 chunks
│ row 2               │   │  → chunks: [0:2], [2:4], [4:5], []
│ row 3               │   │
│ row 4               │  ─┘
│ row 5               │  ─┐
│ row 6               │   ├─ Second piece [5:9], split into 4 chunks
│ row 7               │   │  → chunks: [5:6], [6:7], [7:8], [8:9]
│ row 8               │  ─┘
└─────────────────────┘

Interleaved distribution to ranks:
  Rank 0: rows [0,1] + [5]     = rows [0,1,5]     (3 rows)
  Rank 1: rows [2,3] + [6]     = rows [2,3,6]     (3 rows)
  Rank 2: rows [4]   + [7]     = rows [4,7]       (2 rows)
  Rank 3: []         + [8]     = rows [8]         (1 row)

Step 2: Pad for uniform all-to-all

Before all-to-all, we pad so all ranks have uniform chunk sizes:

• Old dimension (dim 0): max_chunk_size = 3, pad ranks 2 and 3
• New dimension (dim 1): size 4 with 4 chunks → already uniform (chunk size 1 each)

After padding dim 0:
  Rank 0: [0,1,5] (no padding)    → shape (3, 4)
  Rank 1: [2,3,6] (no padding)    → shape (3, 4)
  Rank 2: [4,7,P] (+1 padding)    → shape (3, 4)
  Rank 3: [8,P,P] (+2 padding)    → shape (3, 4)

Step 3: All-to-all on dim 0 → dim 1

The all-to-all exchanges slices: each rank sends dim-1 slices to other ranks and receives dim-0 slices:

Before A2A (each rank has 3x4):     After A2A (each rank has 12x1):
  Rank 0: rows [0,1,5] cols [0,1,2,3]  →  col 0 from all ranks
  Rank 1: rows [2,3,6] cols [0,1,2,3]  →  col 1 from all ranks
  Rank 2: rows [4,7,P] cols [0,1,2,3]  →  col 2 from all ranks
  Rank 3: rows [8,P,P] cols [0,1,2,3]  →  col 3 from all ranks

Step 4: Unpad and reorder

After all-to-all, each rank has interleaved rows from the strided pattern with padding. We use index_select to:

1. Extract only the valid (non-padded) elements
2. Reorder from strided order [0,1,5,2,3,6,4,7,8] back to natural order [0,1,2,3,4,5,6,7,8]

Final result - Shard(1) distribution:
  Rank 0: all 9 rows, col 0  → shape (9, 1)
  Rank 1: all 9 rows, col 1  → shape (9, 1)
  Rank 2: all 9 rows, col 2  → shape (9, 1)
  Rank 3: all 9 rows, col 3  → shape (9, 1)

Pull Request resolved: #170915
Approved by: https://github.com/weifengpy