Device meshes: multiply ranks without multiplying sharding claims

Consider a2×2 mesh whose first axis is data parallel and second is tensor parallel. With ranks arranged[[0,1],[2,3]], tensor-parallel groups are{0,1} and{2,3}, while data-parallel counterpart groups are{0,2} and{1,3}. Each TP group computes one logical model replica on a data shard. Corresponding parameter shards across DP groups need matching initialization and the defined gradient synchronization. Merely creating the mesh does not make a model’s current tensors obey those placements.

Process-group and collective ordering must be consistent across participants. A rank selecting the wrong mesh slice can hang or combine unrelated values. Physical topology also matters for performance: a communication-intensive TP axis can benefit from a stronger interconnect, while the DP axis may have a different traffic pattern. These are placement decisions based on the selected algorithm and hardware, not universal claims that one axis must always be intra-node or that all meshes have equal cost.

For a layer whose weights are evenly tensor-sharded over TP and otherwise replicated in DP, local parameter elements equal that layer’s count divided by TP, not by TP×DP. DP increases the number of copies of each shard. Context parallelism commonly partitions sequence activations while retaining parameter copies across the relevant context ranks; its gradient groups must combine those token contributions. Pipeline parallelism assigns layers to stages, so its parameter accounting depends on which layers a stage actually owns, not a blanket division of every individual layer.

Optimizer-state sharding is another independent placement choice. Adding a sharded optimizer can change local optimizer memory without changing which weight shard a rank needs during the forward pass. Tied parameters across pipeline boundaries need explicit treatment. Distinguish unique logical parameter elements, resident per-rank bytes, and aggregate bytes across the cluster. A chart that divides every category by the total device count can look balanced while badly understating memory needed on each rank.

DP ranks often process different data, whereas TP ranks cooperate on the same logical examples. Averaging a loss across both indiscriminately can introduce extra factors. Sequence/context owners may hold unequal valid-token counts, so global mean scaling needs a valid-count denominator and the framework’s actual sum or average reductions. Track where gradients have already been combined to avoid reducing them twice. A correct mesh is a map of identities and algebraic responsibilities, not just a convenient reshape of rank numbers.

The reliability lab counts available replicas only. Reset, set n=4 and k=0, predict four up with majority threshold 3, then Step. This does not build a device mesh or establish TP/DP membership, and majority has no role in the parameter formulas. It supplies a simple reminder that four available devices are a capacity count, not evidence that every tensor has four distinct shards. The mesh and weight example below carries the actual ownership argument.