Layout redistribution: a shard and a partial sum are different objects

A row shard contains distinct global rows, so reconstructing the full tensor concatenates the owned rows. A sum-partial tensor contains contributions to the same coordinates, so reconstruction sums them. A replica already contains the full logical value. These are not three performance labels for interchangeable buffers. A linear layer’s input-feature products naturally create partial sums, while its output-feature partitions naturally create shards. Carry that placement information through the next operation rather than guessing from a printed local array.

In the documented one-mesh-dimension DTensor transitions, Shard to Replicate uses an all-gather, Partial with sum to Replicate uses an all-reduce, and sum-Partial to Shard uses reduce-scatter. Replicate to Shard can select local chunks without network traffic when replicas are consistent. Changing a shard axis can use all-to-all redistribution. These are version-scoped transition semantics; a larger composed plan can fuse or restructure movement, so its end-to-end communication count still depends on the actual layout graph.

A row-sharded matrix changing to column sharding must send each row owner’s appropriate column pieces to the new owners. Reshaping each local buffer cannot make remote coordinates appear. Conversely, all-reducing distinct row shards would add unrelated rows and destroy their global indexing. The operation needs global shape, shard axis, mesh dimension, and rank order. For uneven partitions, local sizes and offsets also matter; PyTorch 2.8 documents experimental limitations for uneven DTensor sharding, so test the selected operator rather than assuming support.

Converting a DTensor to local tensors can be useful at a kernel boundary, but it transfers responsibility for correct metadata and autograd mapping to the surrounding code. Creating a replicated interpretation from unequal local values does not make them replicas. Creating a sum-partial interpretation from already complete replicas can multiply the value by the mesh size when reduced. Validate both the forward transition and its gradient behavior; a wrong reverse mapping can train a different model even when reconstruction looks right.

Write the global tensor shape, local shape, placement per mesh axis, and intended next placement in a compact table. Record whether a collective completes a sum or only moves coordinates. Check process-group membership and consistent collective order to avoid hangs that a local shape test cannot find. When a benchmark reports bytes, specify whether it counts local payload, sent bytes per rank, received bytes, or aggregate link traffic. Layout labels explain data meaning; a performance model then explains the cost of changing them.

The matrix lab provides a small invertible-transform analogy only. Reset, set a=2 and b=0, predict area scale 2, then Step. The chosen map changes geometric area, whereas a correct layout redistribution preserves the logical tensor’s values. This contrast helps distinguish representation changes from mathematical transformations. The lab does not gather shards, perform reductions, preserve DTensor metadata, or simulate a network; the worked array trace establishes the real redistribution semantics.