Training on Thousands of GPUs, in depth: composing TP, PP and DP, mapping the mesh to the network, and a full step budget for 70B on 1,024 GPUs

Four ways to split a step

There are four ways to divide the work of a transformer training step, and large runs use several at once. Each divides something different and pays for it with a different communication pattern.

The axes multiply: total GPUs = DP x TP x PP x CP (and expert parallelism adds another factor for mixture-of-experts models). The art is choosing the factors. TP cuts weight and activation memory per GPU but talks inside every layer, so its traffic is large and latency-sensitive. PP cuts weight memory with very little traffic but introduces idle time, the pipeline bubble. DP adds throughput with traffic only once per step, but every replica must fit the model, which is why sharded data parallelism (ZeRO stages or FSDP) splits optimizer state, gradients and possibly weights across the DP group as well.

A published reference point: Meta's Llama 3 paper describes training the 405B model with TP=8 and PP=16, using DP=64 on 8,192 GPUs and DP=128 on 16,384 GPUs, at 38 to 43 percent BF16 model FLOPs utilisation, with context parallelism added for the long-context phase. The shape is typical: TP fills a node, PP spans a handful of nodes, DP takes everything that is left.

From factors to process groups

Once the factors are chosen, every rank needs to know which peers it talks to on each axis. The cleanest way to express this is a device mesh: lay the ranks out in an n-dimensional array whose dimensions are the axes, and each process group is a line through that array. With the order (DP, PP, TP) and row-major numbering, a rank's coordinates follow from rank = dp*(PP*TP) + pp*TP + tp, so TP varies fastest. If the launcher assigns ranks contiguously per node, which is the usual behaviour, the eight ranks of each TP group land on the eight GPUs of one server and their traffic never leaves NVLink.

import os, torch, torch.distributed as dist
from torch.distributed.device_mesh import init_device_mesh

TP, PP = 8, 4
world = int(os.environ["WORLD_SIZE"])           # 1024
DP = world // (TP * PP)                          # 32
torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))

mesh = init_device_mesh("cuda", (DP, PP, TP), mesh_dim_names=("dp", "pp", "tp"))
tp_group = mesh["tp"].get_group()    # 8 ranks, same node
pp_group = mesh["pp"].get_group()    # 4 ranks, one per stage, 4 nodes
dp_group = mesh["dp"].get_group()    # 32 ranks, one per replica

rank = dist.get_rank()
dp_i, rem = divmod(rank, PP * TP)
pp_i, tp_i = divmod(rem, TP)
assert mesh["tp"].get_local_rank() == tp_i     # catch launcher surprises early

The assertion at the end is not decoration. If a scheduler hands out ranks round-robin across nodes, TP groups silently span the scale-out network, every layer's all-reduce crosses InfiniBand or Ethernet instead of NVLink, and the job runs several times slower with no error at all. Check the mapping at start-up and refuse to train if it is wrong.

Mapping the mesh onto the network

The ordering principle is simple: put the axis with the most frequent, most latency-sensitive traffic on the fastest links. Inside an H100 server, NVLink gives each GPU 900 GB/s of aggregate bandwidth. Between servers, a typical design gives each GPU one 400 Gb/s network port, about 50 GB/s, an order of magnitude less. So TP goes inside the node; PP, whose messages are small and few, can cross nodes freely; DP traffic, which is large but happens once per step and overlaps backward, rides the scale-out fabric.

Two second-order effects matter at thousands of GPUs. First, DP collectives across many racks hit the upper tiers of the network, so the scheduler should place each pipeline's nodes, and ideally whole replicas, under the same leaf switches where it can. Second, the rail-optimised wiring common in GPU clusters connects GPU k of every server to the same leaf switch; DP groups that consist of the same local GPU index on every node (exactly what the mesh above produces) then talk mostly within a rail. NCCL discovers the topology itself, but it cannot fix a bad rank layout.

Anatomy of one step

Within one optimizer step, each pipeline processes its share of the global batch as m microbatches. With the common one-forward-one-backward (1F1B) schedule, a stage runs a warm-up of forwards, then alternates forward and backward, then drains backwards. The idle time at the start and end is the bubble, a fraction (p - 1) / (m + p - 1) of the step for p stages. More microbatches shrink it; interleaved schedules that give each GPU several non-adjacent chunks of layers shrink it further at the cost of more point-to-point messages.

# One training step on one rank (1F1B, simplified)
for mb in warmup_microbatches(pp_i, PP, m):
    act = recv_forward(pp_group)            # from stage pp_i - 1 (stage 0 reads data)
    out = stage_forward(act)                # TP all-reduces happen inside, on NVLink
    send_forward(out, pp_group)
for mb in steady_state(m):
    out = stage_forward(recv_forward(pp_group)); send_forward(out, pp_group)
    grad = stage_backward(recv_backward(pp_group)); send_backward(grad, pp_group)
    # DP gradient buckets for finished layers start reducing here, overlapped
for mb in cooldown(pp_i, PP):
    send_backward(stage_backward(recv_backward(pp_group)), pp_group)
finish_dp_reduce(dp_group)                  # wait for the last buckets
clip_grad_norm_across(mesh)                 # norm needs a reduction over tp, pp and dp
optimizer.step(); optimizer.zero_grad()

Note the gradient clipping line. The global gradient norm is a sum over every parameter shard on every axis, so it needs its own small all-reduce. Computing it only over the local shard is a classic silent bug: training proceeds, but clipping uses a different threshold on every rank.

Worked example: 70B on 1,024 GPUs

Take a dense 70B model with 80 layers, hidden size 8,192 and sequence length 4,096, trained on 1,024 H100s (128 servers) with TP=8, PP=4 and therefore DP=32. The global batch is 1,024 sequences, about 4.19 million tokens; each replica handles 32 sequences per step as m = 32 microbatches of one sequence.

Read the table as a budget. Memory is comfortable: under 10 GB of weights and gradients plus under 1 GB of sharded optimizer state leaves most of an 80 GB GPU for activations, which is what lets each microbatch hold a whole 4,096-token sequence. The DP traffic is small next to a 4.35 s step and hides behind backward easily. The real pressure is TP: roughly 0.67 s of NVLink traffic per step, about 15 percent of the step, sits on the critical path unless it is overlapped with the matrix multiplies or reduced by sequence parallelism (which replaces the all-reduce with a reduce-scatter and an all-gather and shards the layer-norm activations). Add the 8.6 percent bubble and you can see why 40 percent MFU is a good outcome rather than a pessimistic one.

With the global batch fixed, m = 1,024 / DP, so every change moves m too. PP=8 gives DP=16 and m=64: weights per GPU halve and the bubble rises only to 7/71, about 9.9 percent. TP=4 gives DP=64 and m=16: TP traffic falls to about 0.29 s per step, but weights per GPU double and the bubble grows to 3/19, about 16 percent. A larger global batch improves every ratio but changes optimisation, so it is a research decision.

Running it: start-up, hangs, stragglers, data

Getting a thousand processes to start, agree and keep agreeing is its own discipline.

• Start-up. Rendezvous (for example torchrun with a shared store) and NCCL communicator creation for every group take noticeable time at this scale. Create groups once, eagerly, in the same order on every rank; a rank that creates groups in a different order deadlocks the job at initialisation.
• Timeouts and hang forensics. A collective waits forever for a missing peer, so set the process-group timeout deliberately and turn on PyTorch's NCCL flight recorder (TORCH_NCCL_TRACE_BUFFER_SIZE) so a hang leaves a record of which collective each rank was in. The rank that is not in the collective is your suspect.
• Stragglers. A synchronous job runs at the speed of its slowest GPU. Log per-rank forward, backward and optimizer times every N steps, gather them, and alert on ranks persistently slower than the median. Thermal throttling, a degraded link and a noisy neighbour on shared storage all show up this way first.
• Data. Each DP replica must read a disjoint shard, while every rank in a TP group must see identical inputs. Make the data loader's position part of the checkpoint so a restart neither repeats nor skips tokens.
• Numerics. Loss spikes are more frequent and costlier at scale. Keep the ability to skip a batch or rewind to the last good checkpoint without manual surgery, and log gradient norm per step.

Failures are certain at this scale. How often they occur and how to survive them are treated in GPU hardware faults; the companion cluster design article derives the fabric and storage that this layout assumes.

Failure modes

Trade-offs

No layout is best in general; each axis trades memory against communication and idle time. Prefer the smallest TP that fits the model and keeps matrix multiplies large enough to be efficient, because TP traffic grows with every layer. Use PP when weights still do not fit, and buy back the bubble with more microbatches. Use sharded data parallelism before adding pipeline stages when optimizer state, not weights, is what overflows. Add context parallelism only when sequence length forces it. A simpler layout is also easier to debug, which counts for a lot during a months-long run.

What to do next

1. Write down the model's parameter count, layer count, hidden size and sequence length, and reproduce the budget table above for your own numbers before choosing a layout.
2. Pick TP to fit one NVLink domain, then PP, then let DP absorb the remaining GPUs.
3. Build the device mesh in one function and assert that every TP group is on one node.
4. Read tensor parallelism, pipeline parallelism and the ZeRO optimizer for the mechanics of each axis.
5. Profile one step at small scale (for example 32 GPUs with the same TP and PP) and confirm the bubble and TP communication match your estimate.
6. Turn on per-rank step timing and the NCCL flight recorder before the first large launch.
7. Choose the checkpoint interval with training checkpointing before the first large launch, using your measured step time.