Making Training Fast

4.1Warp & SIMT

NVIDIA GPUs execute threads in groups of 32 called warps. This is Single Instruction, Multiple Threads (SIMT): the warp shares one instruction stream, while each thread has its own registers.

Two consequences dominate kernel performance:

1. Branch divergence: if threads in a warp take different paths of an if, the warp runs both paths with lanes masked off. Keep control flow uniform within a warp where you can.
2. Memory coalescing: when the 32 threads read adjacent addresses, the hardware merges them into a few wide transactions. Scattered access multiplies memory traffic. See Memory coalescing.

Warps are also the unit the scheduler juggles to hide latency. More resident warps (higher Occupancy) gives it more choices, up to the point where registers and shared memory run out.

4.2Memory coalescing

Global memory is accessed in aligned segments (32-byte sectors, grouped into 128-byte lines). If the 32 threads of a Warp & SIMT read 32 consecutive 4-byte floats, the hardware serves them with one 128-byte transaction. If they read with a large stride, it may need 32 separate transactions, a 32× bandwidth penalty.

Classic fixes from GPU programming practice (e.g. the PRACE guide)

• Structure-of-arrays instead of array-of-structures.
• Have thread i touch element i, not i × stride.
• Stage awkward access patterns (transposes, stencils) through shared memory, then write out coalesced.
• Watch for shared-memory bank conflicts when you do.

Profilers (Nsight Compute) report "sectors per request". Ideal is 4 for 4-byte loads.

4.3Occupancy

Each SM can host a limited number of warps (64 on H100). Occupancy is how many are actually resident, limited by each block's use of:

• Registers per thread (the register file is shared by all resident warps),
• Shared memory per block,
• Threads per block and blocks per SM.

Higher occupancy gives the warp scheduler more candidates to switch to while others wait on memory. But it's a means, not a goal. Many fast kernels (big GEMMs, FlashAttention) run at low occupancy, using registers and shared memory for data reuse instead.

Practical approach (from CUDA and PRACE best practice): use the occupancy calculator or Nsight Compute to see the limiting resource, try a few block sizes (multiples of 32), and measure. Don't assume.

4.4Tensor Core

Tensor cores compute small matrix products (e.g. 16×16 tiles) in one instruction. On an H100 they deliver ~989 dense BF16 TFLOP/s versus ~67 TFLOP/s of plain FP32. That gap is why Mixed precision (BF16, FP8) training is the default.

Getting the FLOPs

• Use BF16/FP16 (or FP8 with Transformer Engine) for matmuls. FP32 without TF32 never touches them.
• Keep GEMM dimensions large and multiples of 8/16. Tiny micro-batches or very narrow tensor-parallel shards make small GEMMs that can't saturate the units.
• Keep data flowing from HBM & the memory hierarchy. Tensor cores are so fast that many ops (normalization, activations, attention softmax) become memory-bound. Kernel fusion such as FlashAttention exists for this.

Model FLOPs Utilization (MFU (Model FLOPs Utilization)) is ultimately a measure of how busy the tensor cores are.

4.5Mixed precision (BF16, FP8)

Mixed-precision training runs forward and backward in a low-precision format and keeps FP32 master weights for the optimizer update:

• BF16 keeps FP32's exponent range, so it rarely needs loss scaling. It's the default on A100/H100.
• FP16 needs dynamic loss scaling to keep small gradients from underflowing.
• FP8 (H100+) roughly doubles matmul throughput again. Libraries like Transformer Engine manage per-tensor scaling.

Why it's a big deal: Tensor cores deliver ~15× more FLOP/s in BF16 than plain FP32 on H100. Training in FP32 leaves almost all of the GPU idle.

Memory: weights and gradients in 2 bytes, plus 12 bytes of FP32 optimizer state (master, Adam m, v) gives the "16 bytes per parameter" rule used throughout the Atlas.

4.6Data parallelism (DDP)

The simplest way to scale: replicate the model on N GPUs, give each a different slice of the batch, then all-reduce the gradients so every replica takes the same optimizer step.

Cost: one all-reduce of the full gradient (2 bytes × params in BF16) per step. With ring all-reduce the per-GPU traffic is ~2× the gradient size, independent of N. That's why data parallelism scales so well.

Making it fast

• Overlap: DDP buckets gradients and all-reduces each bucket as soon as backward produces it, hiding communication behind compute (see CUDA streams & async execution).
• Gradient accumulation: run several micro-batches before syncing. You get a bigger global batch with the same memory and fewer syncs per token.

Limit: every GPU must hold the whole model state. Once that stops fitting, see ZeRO & FSDP, Tensor parallelism and Pipeline parallelism.

4.7CUDA streams & async execution

A stream is an ordered queue of GPU operations. Work in different streams can run concurrently: a host-to-device copy in one, a kernel in another, an NCCL collective in a third. Events express dependencies between streams without blocking the CPU.

This is the mechanism behind every "overlap" optimization:

• Gradient all-reduce overlapping with the backward pass in DDP uses a separate communication stream.
• Prefetching the next batch while the current one computes (see Input pipeline).
• ZeRO-3/FSDP prefetching the next layer's parameters (see ZeRO & FSDP).

Projects such as grCUDA go further. They build the dependency graph from the data each kernel touches and schedule independent kernels on separate streams automatically, so polyglot programs get concurrency without hand-written stream code.

Pitfall: any implicit synchronization (cudaDeviceSynchronize, .item() in PyTorch, pageable memory copies) collapses your carefully built overlap back into serial execution.

4.8Activation checkpointing

Backpropagation needs the activations from the forward pass. For a transformer layer with FlashAttention that's roughly 34 · s · b · h bytes (sequence × micro-batch × hidden), and it adds up fast. A 7B model at micro-batch 4 needs over 70 GB of activations.

Checkpointing (also called recomputation) keeps only each layer's input, then re-runs the layer's forward during backward. Cost: an extra forward pass, ~33% more compute (8N instead of 6N FLOPs per token). Benefit: activation memory drops to about 2 · s · b · h bytes per layer.

Use it when activations are what's stopping you from fitting the model or a larger micro-batch. Turn it off when memory is fine. It's pure overhead then. Selective checkpointing (only the attention internals) gets most of the savings for a fraction of the cost.

Note that MFU deliberately doesn't count recomputed FLOPs, so checkpointing lowers MFU even though the GPU is just as busy.

4.9ZeRO & FSDP

Plain Data parallelism (DDP) wastes memory: every rank stores identical copies of everything. ZeRO (and PyTorch's FSDP) shards it:

The catch with stage 3: parameters must be all-gathered before each layer's forward and backward, every micro-batch. For a 70B model that's ~140 GB moved twice per micro-batch. On NVLink that's fine. Over a 400 Gb/s NIC it can dominate the step.

Practical guidance

• ZeRO-1 is nearly free and almost always worth it.
• ZeRO-3/FSDP is excellent within a node or on fast fabrics, and for fitting a model that otherwise wouldn't fit.
• At large scale, mix it: TP inside the node, ZeRO-1 or hybrid sharding (shard within a node, replicate across nodes).

4.10Tensor parallelism

Tensor parallelism (TP) shards the weight matrices inside each layer. In the Megatron-LM scheme, the first MLP matmul is split by columns and the second by rows, so a layer needs only one all-reduce in forward and one in backward. With attention handled the same way, that's ~4 activation all-reduces per transformer layer per micro-batch.

Trade-offs

• ✅ Weights, gradients, optimizer state and activations shrink by the TP degree.
• ❌ Communication is frequent, latency-sensitive and on the critical path. It mostly can't be hidden.
• ❌ Shards get narrow: hidden/TP too small means small GEMMs and lower Tensor Core efficiency.

Rule: keep TP inside the NVLink domain (TP ≤ 8 on HGX servers). TP=16 across InfiniBand & RoCE is a classic way to waste a cluster.

4.11Pipeline parallelism

Pipeline parallelism (PP) assigns contiguous layers to stages. Stage 1 runs layers 1–20, stage 2 runs 21–40, and so on. Only the activations at stage boundaries cross the network, so PP is comfortable spanning nodes over InfiniBand & RoCE.

The bubble: at the start and end of each step some stages sit idle while the pipeline fills and drains. With p stages and m micro-batches per step (1F1B schedule):

bubble fraction ≈ (p − 1) / (m + p − 1)

More micro-batches (gradient accumulation) shrink the bubble. Interleaved schedules shrink it further, at the cost of more communication.

Constraints: the layer count must divide evenly (or nearly) into stages, and stages should be balanced. The first stage also holds the most in-flight activations.

The classic large-model recipe, 3D parallelism, combines all three: TP inside the node, PP across a few nodes, DP (often with ZeRO-1) across everything else.