Distributed training

What does DistributedDataParallel do each step?

It buckets gradients and allreduces them, overlapping the communication with the backward of later layers. The weights end up identical on every rank.

What memory does ZeRO-1 remove from each GPU?

The optimizer state, which for Adam is about twice the parameter memory in fp32 moments. Weights and gradients stay replicated.

What does ZeRO-2 add to stage 1?

Sharded gradients as well as optimizer state. Parameters are still gathered as a full copy for the forward.

What does ZeRO-3 gather, and when?

Parameters for the layer that is about to run, then it discards them. The step does an allgather on the way in and a reduce-scatter on the way out.

How does FSDP map onto that idea?

It is PyTorch's shard of parameters, gradients, and optimizer state, with a wrap unit you choose, often a transformer block. The gather happens at that unit's boundary.

Why does the FSDP wrap size matter?

Wrapping each linear layer gathers too often. Wrapping the whole model gathers almost everything at once. A block is the usual middle.

What is tensor parallelism?

One layer's matrix is split across GPUs. Each GPU computes a shard and they exchange activations inside the layer, so that link is on the step's critical path.

When does tensor parallelism communicate?

On every forward and backward of the split layers, not once per step. The GPUs need a fast link. Across a rack it often loses to another split.

What is sequence parallelism beside tensor parallelism?

Sequence activations are split across the tensor-parallel GPUs, so layernorm and dropout do not keep a full copy. It is the partner of Megatron tensor parallel.

What is pipeline parallelism?

Groups of layers live on different GPUs. A microbatch moves from one stage to the next. The weight of the model is split by depth.

What is a pipeline bubble?

The idle time while the first microbatch fills the pipeline and the last one drains it. More microbatches hide a larger fraction of that idle time.

What is a microbatch?

The slice of the batch that one pipeline stage holds at a time. Gradient accumulation over microbatches rebuilds the logical batch.

What is 1F1B pipeline scheduling?

A stage interleaves one forward and one backward so it does not have to store activations for every microbatch. Activation memory stays bounded by a few microbatches.

What is 3D parallelism?

Data, tensor, and pipeline parallel together. Tensor parallel stays inside a fast node. Pipeline splits depth. Data parallel replicates that pair across groups.

What is context parallelism?

The sequence itself is split across GPUs for long contexts, and attention exchanges the KV it needs. It is how a 2026 long-context run fits the activation memory.

What is expert parallelism?

In a mixture of experts, different experts live on different GPUs. Tokens are routed to the GPU that holds their expert, then gathered back.

What fails if expert routing is unbalanced?

One GPU does most of the expert compute and the others wait. A balancing loss or a capacity cap keeps the step from following the hottest expert.

What does Megatron-Core implement?

The tensor, pipeline, context, and expert splits for a transformer, plus the communication they require. You use it when data-parallel sharding is not enough.

How does JAX express the same split?

You shard named axes of the arrays and the compiler inserts the collectives. The training step stays a function. The mesh says which device owns which axis.

What is a device mesh?

A grid of GPUs with named axes, such as data and tensor. A tensor's sharding is a placement on those axes. Changing the mesh changes the communication.

Why is bf16 the usual 2026 compute type?

It has fp32's exponent range, so activations overflow less than in fp16, and the GPU tensor cores run it. Reductions and the master weights often stay fp32.

What does FP8 training require?

Hardware with FP8 tensor cores, and a scaling scheme so the small values survive. You measure loss against the bf16 run. The speedup is not free of a quality check.

What are master weights?

An fp32 copy that the optimizer updates. The bf16 or fp8 weights are a cast used for the matmul. Updating only the low-precision copy drifts.

Why is Adam the memory problem?

Each parameter carries two fp32 moments. That is more bytes than the bf16 weight. Sharding or an 8-bit optimizer targets those moments first.

What does activation checkpointing store?

The inputs of chosen blocks, not every intermediate. The backward recomputes the rest. Memory falls and the step does extra compute.

What is selective activation checkpointing?

Only the most expensive activations are dropped, often the attention scores, so you recompute less than a full block checkpoint and still fit.

What is model FLOPs utilization?

Achieved FLOPs per second divided by the GPU's peak. A low MFU with a small model usually means you are latency-bound on kernels, not short of math.

When is a training step compute-bound?

When the matmul is large enough that communication and kernel launch are hidden behind it. Large hidden sizes and long sequences push that way.

When is a step communication-bound?

When the allreduce, allgather, or pipeline send takes longer than the math it overlaps. A small model on many GPUs often lands here.

What does overlapping communication mean?

The gradient allreduce of an earlier layer runs while a later layer is still in backward. If the math is shorter than the network transfer, overlap cannot hide all of it.

What is a gradient bucket?

A group of gradients allreduced together so the network sees large messages. Buckets that are too small waste the link. Buckets that are too large delay the overlap.

What do you read in the GPU topology before you pick a split?

Which GPUs share NVLink and which hops cross the network. Tensor parallel wants the NVLink group. Data parallel can cross the slower link.

What is a straggler in data parallel?

The rank whose batch or machine is slowest. Every other rank waits at the collective. One slow data loader stalls the whole step.

A loss spike hits one rank's batch. What do you check?

That rank's data, the global gradient norm, and whether a collective skipped or overflowed. A spike on every rank at the same step is a learning-rate or data-order event.

Why log the global gradient norm?

It is the scale of the update before clipping. A sudden jump often precedes a loss spike. Clipping hides the jump in the weights and leaves it visible in the log.

What does gradient clipping change?

The direction is kept and the length is capped. It stops one step from jumping. It does not fix a loss that is computed on the wrong tokens.

What does gradient accumulation simulate?

A larger batch, by summing gradients over several forwards before one optimizer step. The optimizer and the allreduce see the accumulated gradient.

How does accumulation interact with ZeRO?

You still accumulate locally, then the sharded reduce uses the summed gradient. The effective batch is microbatches times accumulation times data-parallel ranks.

What is the critical batch size?

The batch beyond which tokens per second keep rising but the loss improvement per token falls off. Past it you spend GPUs on almost the same learning.

Why count tokens and not steps?

A step's size changes with packing, accumulation, and sequence length. A paper or a run is compared at the same number of tokens seen.

What does a resumable checkpoint contain?

Weights, optimizer moments, the data-loader position, and the step count, from the same step. Weights alone restart the learning rate and repeat data.

Why is a ZeRO or FSDP checkpoint sharded?

Each rank writes its shard, because the full optimizer state may not fit on one disk or one process. Loading requires the same shard count, or a merge.

Why can the same seed still diverge across runs?

Atomic adds in a reduction, a different allreduce order, and nondeterministic kernels. Bit-exact reruns need deterministic flags and the same GPU count.

What does compiling the training step risk?

A long first compile, and a recompile when shapes change. The speedup shows up after that, on the shapes you actually train.

What does FlashAttention change in a training step?

The attention math result stays the same, within numeric tolerance, and the memory traffic drops. The backward is fused too. It is a kernel swap, not a new objective.

Why do research codebases vendor Triton kernels?

A fused kernel for custom attention or MoE routing beats the eager PyTorch op. Test it against a slow reference. The kernel is part of the experiment.

What is NCCL doing in the step?

The GPU-to-GPU collectives: allreduce, allgather, reduce-scatter. A hang at a collective is a rank that never arrived, not a Python exception on the fast ranks.

One rank OOMs and the others sit silent. Why?

They are blocked in a collective waiting for the dead rank. The failure looks like a hang. Timeouts and a process group that aborts together make it a crash you can see.

What does a distributed run report so it can be compared?

GPU type and count, the parallel sizes, tokens per step, MFU, achieved tokens per second, and the loss at a stated token count.