Mesh, placement, and the collective it implies

init_device_mesh("cpu", (2, 2), mesh_dim_names=("dp", "tp")) on a world of four builds a 2 by 2 grid and gives its axes names. Print it and you get DeviceMesh([[0, 1], [2, 3]]): ranks 0 and 1 sit along the tp axis of the first dp row, ranks 2 and 3 along the second.

The names are the working part. Once an axis is called dp and another tp, a layout stops being a tuple of integers you have to keep straight and becomes a statement about which parallelism runs where. A subgroup for either axis comes out of the mesh rather than out of a hand-built new_group call, which is the same coordination lesson one showed, one level up.

This grid is the same object the jax path's sharding chapter builds with a mesh and a PartitionSpec, and the same grid the xla path's SPMD chapter says the partitioner rewrites a program against. Three courses, one idea, arrived at from three directions. The reason to state that here rather than teach it three times is that the mechanism below is where PyTorch's version differs.

from torch.distributed.device_mesh import init_device_mesh

mesh = init_device_mesh("cpu", (2, 2), mesh_dim_names=("dp", "tp"))
if rank == 0:
    print("mesh:", mesh)
    print("mesh shape:", mesh.shape, "dim names:", mesh.mesh_dim_names)

# mesh: DeviceMesh([[0, 1], [2, 3]])
# mesh shape: (2, 2) dim names: ("dp", "tp")

A DTensor is a global tensor plus a list of placements, one entry per mesh axis, each saying how that axis divides the tensor. Shard(0) on the first axis and Shard(1) on the second means: split rows across the dp axis, split columns across the tp axis.

Distribute a 4 by 4 tensor of 0 through 15 that way and the block each rank holds is exactly the one the two placements imply. Rank 0 gets [0, 1, 4, 5], the top-left 2 by 2. Rank 3 gets [10, 11, 14, 15], the bottom-right. The DTensor's .shape still reports (4, 4) on every rank while .to_local().shape reports (2, 2), which is the difference between what the program is computing with and what this process is holding.

The third placement type has no equivalent in a shape. Replicate() says an axis does not divide the tensor at all, so every rank along it holds the same values. _Partial says something stranger, and the section after next is where it earns its own paragraph.

from torch.distributed._tensor import Shard, distribute_tensor

big = torch.arange(16.).reshape(4, 4)
d = distribute_tensor(big, mesh, [Shard(0), Shard(1)])
print(f"rank {rank}: global {tuple(d.shape)} local {tuple(d.to_local().shape)} "
      f"{d.to_local().flatten().tolist()}")

# rank 0: global (4, 4) local (2, 2) [0.0, 1.0, 4.0, 5.0]
# rank 1: global (4, 4) local (2, 2) [2.0, 3.0, 6.0, 7.0]
# rank 2: global (4, 4) local (2, 2) [8.0, 9.0, 12.0, 13.0]
# rank 3: global (4, 4) local (2, 2) [10.0, 11.0, 14.0, 15.0]
# placements: (Shard(dim=0), Shard(dim=1))

Ask a DTensor for a different layout and it works out what has to move. CommDebugMode is the instrument that makes the answer visible: a dispatch mode that counts functional collectives inside its context, shipped in torch.distributed._tensor.debug at 2.2.2 for exactly this purpose.

Four transitions, on a one-dimensional mesh of four, and the counts are not symmetric. Going from Shard(0) to Replicate() costs one all-gather, because every rank has to receive what the others hold. Going the other way, Replicate() to Shard(0), costs nothing at all: every rank already has every value, so it just keeps its slice and drops the rest.

That asymmetry is the whole reason to think in placements rather than in calls. Sharding an already-replicated tensor is free. Replicating a sharded one is a collective. The gym's naming drill at GYM·11 works the same skill from the other end, giving you the before-and-after tensors and asking which collective produced them; this section gives you the placements and asks which collective they imply.

Sharding a replicated tensor is free. Replicating a sharded one is an all-gather.

Shard(0) -> Replicate : {"all_gather_into_tensor": 1}
Replicate -> Shard(0) : {}
Shard(0) -> Shard(1)  : {"all_gather_into_tensor": 1}
Shard(1) @ Shard(0)   : {} -> placements (_Partial(reduce_op=RedOpType.SUM),)
Replicate @ Shard(1)  : {} -> placements (Shard(dim=1),)
full_tensor of the matmul: {"all_reduce": 1}
result row 0: [6.0, 6.0, 6.0, 6.0]

import os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.distributed._tensor import Replicate, Shard, distribute_tensor
from torch.distributed._tensor.debug import CommDebugMode
from torch.distributed.device_mesh import init_device_mesh

def counts(mode):
    return {str(k).replace("c10d_functional.", ""): v
            for k, v in mode.get_comm_counts().items()}

def worker(rank, world):
    os.environ["MASTER_ADDR"] = "127.0.0.1"
    os.environ["MASTER_PORT"] = "29542"
    dist.init_process_group("gloo", rank=rank, world_size=world)
    mesh = init_device_mesh("cpu", (4,))
    x = torch.arange(16.).reshape(4, 4)
    w = torch.ones(4, 4)

    sharded = distribute_tensor(x, mesh, [Shard(0)])
    m = CommDebugMode()
    with m:
        r = sharded.redistribute(mesh, [Replicate()])
    if rank == 0:
        print("Shard(0) -> Replicate :", counts(m))

    m2 = CommDebugMode()
    with m2:
        s = r.redistribute(mesh, [Shard(0)])
    if rank == 0:
        print("Replicate -> Shard(0) :", counts(m2))

    m3 = CommDebugMode()
    with m3:
        s.redistribute(mesh, [Shard(1)])
    if rank == 0:
        print("Shard(0) -> Shard(1)  :", counts(m3))

    xs = distribute_tensor(x, mesh, [Shard(1)])
    ws = distribute_tensor(w, mesh, [Shard(0)])
    m4 = CommDebugMode()
    with m4:
        out = torch.matmul(xs, ws)
    if rank == 0:
        print("Shard(1) @ Shard(0)   :", counts(m4), "-> placements", out.placements)

    m6 = CommDebugMode()
    with m6:
        full = out.full_tensor()
    if rank == 0:
        print("full_tensor of the matmul:", counts(m6))
        print("result row 0:", full[0].tolist())
    dist.destroy_process_group()

if __name__ == "__main__":
    mp.spawn(worker, args=(4,), nprocs=4, join=True)

Multiply two DTensors whose sharding does not line up and something unexpected happens: no collective runs. A tensor sharded on its columns times a weight sharded on its rows is a valid local matmul on every rank, and each rank's answer is a partial sum, correct in shape and incomplete in value. The result comes back with placement _Partial(reduce_op=RedOpType.SUM), and the count is zero.

The reduction is not skipped, it is owed. Call full_tensor() on that result and one all_reduce fires, and the value comes out right: row 0 of the answer is [6.0, 6.0, 6.0, 6.0], which is 0 + 1 + 2 + 3 against a weight of ones. Deferring the reduction is what lets a chain of operations run without paying for a collective between every pair of them.

The same tracking runs the other way. A replicated input times a column-sharded weight needs no communication and the answer is already Shard(1), ready for a row-sharded layer to consume. Placement algebra is doing the bookkeeping a hand-written distributed program would do in comments.

Chapter 8 sets the goal of predicting which collective a given call inserts. This is the mechanism that makes it predictable: the placements of the inputs determine the placement of the output, and a collective appears only where a placement has to change.

The tensor-parallel API is a thin layer over everything above. Hand parallelize_module a mesh and a dictionary saying ColwiseParallel for the up-projection and RowwiseParallel for the down-projection, and it replaces both weights with DTensors carrying the placements those names describe.

Print them and the naming makes sense: the colwise layer's weight is Shard(0) on a (8, 4) weight, which splits the output features, and the rowwise layer's is Shard(1) on a (4, 8) weight, which splits the input features. Both ranks hold a local (4, 4). The two shardings are chosen to fit together, so the intermediate activation stays sharded across the ReLU and never needs gathering.

One collective runs in the whole forward: a single all_reduce, at the end, where the row-sharded output's partial sums have to be combined. That is the classic tensor-parallel MLP, and here it is a count rather than a diagram.

The output comes back as an AsyncCollectiveTensor rather than a plain tensor, which is the wrapper that lets the all_reduce be issued now and waited on at first use. The xla path's collectives chapter describes the same split as an async start and done pair in the compiled schedule. Both are the same idea: launch the communication, keep computing, wait only when the value is genuinely needed.

from torch.distributed.tensor.parallel import (
    ColwiseParallel, RowwiseParallel, parallelize_module)

mesh = init_device_mesh("cpu", (2,), mesh_dim_names=("tp",))
tp = parallelize_module(MLP(), mesh,
                        {"up": ColwiseParallel(), "down": RowwiseParallel()})
mode = CommDebugMode()
with mode:
    out = tp(torch.ones(2, 4))

# up.weight  : DTensor (8, 4) (Shard(dim=0),)
# down.weight: DTensor (4, 8) (Shard(dim=1),)
# local shapes: (4, 4) (4, 4)
# forward collectives: {"all_reduce": 1}
# out: AsyncCollectiveTensor (2, 4)

Two version facts are worth carrying out of this lesson. At 2.2.2 the tensor type lives at torch.distributed._tensor, with the leading underscore that means the API is not promised to hold still, and CommDebugMode says in its own docstring that it counts functional collectives only, so the FSDP counts from the previous lesson had to be taken by wrapping dist functions instead. Both of those have moved since. Check the import path against your own build before copying a line.

What has not moved is the shape of the idea. A named grid of devices, a per-axis statement of how each tensor sits on it, and a rule that derives the communication from a change in that statement. The jax path teaches it as a mesh plus a PartitionSpec handed to a compiler, and the xla path teaches what the compiler then does to the module. PyTorch does it eagerly, one operation at a time, which is why you can count the collectives with a dispatch mode instead of reading a compiled schedule.

The gym station at GYM·11 is the practice half of this. It shows four ranks' tensors before and after one collective, taken from a real four-process gloo run, and asks which collective ran. Do a streak of five there once the placement transitions above feel obvious, because naming a collective from its effect and predicting one from a placement change are the same knowledge tested from two sides.