What Actually Happens During One Distributed Training Step?

The familiar training loop makes a distributed step look sequential: load a batch, call the model, compute a loss, run backward(), then update the weights. The machine executes something more interesting. CPU workers prepare future batches while copy engines move the current one, GPU kernels build activations, the autograd engine walks the graph in reverse, and NCCL exchanges gradient buckets as soon as they become ready.

for inputs, targets in loader:
    optimizer.zero_grad(set_to_none=True)
    outputs = model(inputs)
    loss = criterion(outputs, targets)
    loss.backward()
    optimizer.step()

Those five lines coordinate several independent schedulers. Python schedules operators. The CUDA runtime schedules work into streams. The GPU schedules thread blocks onto streaming multiprocessors. PyTorch autograd schedules backward operations from dependency edges. DistributedDataParallel (DDP) schedules collectives when gradient buckets become ready. NCCL schedules transfers over NVLink, PCIe, or the cluster fabric.

Understanding the boundaries between those schedulers explains many otherwise surprising observations: why Python can reach the next line before a kernel finishes, why gradient communication starts before backward() returns, why one slow rank delays every rank, and why the optimizer usually runs independently on every GPU rather than on a central server.

The step is a pipeline, not a list

A useful first approximation is to picture four lanes:

1. The input lane reads, decodes, tokenizes, augments, and batches samples on CPUs.
2. The transfer lane copies a prepared batch from host memory into GPU memory.
3. The compute lane runs forward kernels, loss kernels, backward kernels, and optimizer kernels.
4. The communication lane reduces gradients across ranks.

A well-fed step overlaps these lanes. While the GPU computes on batch n, workers prepare batch n + 1. During backward, communication for an early gradient bucket can overlap with computation of later gradients. The step time is therefore not the sum of every phase. It approaches the length of the exposed critical path: the work that could not hide behind something else.

1. A rank receives its own batch

In ordinary data-parallel training, each process controls one GPU and is called a rank. Every rank owns a replica of the model. A distributed sampler assigns different samples to different ranks so they do not all train on the same batch.

Before the loop body begins, a PyTorch DataLoader may already have several batches in flight. With num_workers > 0, worker processes fetch samples and run dataset transforms outside the training process. The main process receives completed batches through multiprocessing queues. With pin_memory=True, PyTorch moves recognized tensors into page-locked host memory, from which the CUDA driver can transfer them efficiently.

The rank then requests a host-to-device copy, commonly with:

inputs = inputs.to(device, non_blocking=True)
targets = targets.to(device, non_blocking=True)

non_blocking=True does not mean the data is already on the GPU when to() returns. It means the host can continue after enqueueing the transfer. CUDA stream ordering ensures that a later kernel in the same stream does not consume the batch before its copy completes. Useful overlap with computation requires pinned source memory, compatible hardware, and a schedule that places the copy where it can run concurrently.

2. The forward pass creates both outputs and obligations

Calling model(inputs) enters Python module code, but most expensive tensor operations quickly dispatch CUDA kernels from libraries such as cuBLAS, cuDNN, and custom fused kernels. Matrix multiplications, normalizations, attention, elementwise operations, and reductions become work queued to CUDA streams.

The forward pass produces two kinds of data:

• Visible outputs, such as logits passed to the loss function.
• Saved tensors needed later to compute derivatives, such as layer inputs, normalization statistics, masks, and attention intermediates.

PyTorch autograd records how output tensors depend on earlier operations. It does not store a second copy of the model's Python code. It builds a dynamic graph of backward functions and saves the tensor values those functions will need. This is why activation memory grows during the forward pass and why activation checkpointing trades memory for recomputation: it deliberately discards selected saved tensors and recreates them during backward.

The loss function is part of the same device program. Cross-entropy is not one magical instruction; it typically becomes several kernels or a fused implementation that performs normalization, indexing, and reduction. The resulting loss may be a one-element CUDA tensor. Calling loss.item() asks the CPU for that value and therefore creates a synchronization point: Python cannot receive the scalar until the GPU has produced it.

3. Backward walks dependencies in reverse

loss.backward() seeds the loss gradient and asks the autograd engine to execute the graph in reverse dependency order. A layer near the output can begin its backward work before a layer near the input. Each backward function consumes an incoming gradient, reads any saved tensors, and produces gradients for its inputs and parameters.

Backward therefore has its own memory and execution shape:

• Saved forward activations become eligible for release after their backward consumer finishes.
• Parameter gradients appear progressively, usually in the reverse of forward parameter use.
• Temporary workspaces come and go as backward kernels execute.
• Activation-checkpointed regions rerun part of the forward pass before computing their gradients.

A gradient is not normally sent to a central coordinator. Every rank computes gradients from its local batch. DDP then makes those gradients identical across ranks before each rank applies its own optimizer update.

4. DDP reduces buckets, not a finished model-sized gradient

DDP registers autograd hooks on parameters. When autograd finishes a parameter's gradient, its hook marks that gradient ready in DDP's reducer. Parameters are grouped into contiguous buckets; PyTorch's documented default bucket cap is 25 MiB unless the application changes it.

When every gradient in a bucket is ready, DDP can launch that bucket's all-reduce immediately. It does not need to wait for the complete backward pass. This is the main source of compute/communication overlap in ordinary DDP:

An all-reduce combines corresponding bucket values from every rank and leaves every rank with the same result. DDP averages the gradients across the process group. If eight ranks each used a local batch of 32 samples, the synchronized gradient usually represents a global batch of 256 samples, subject to the loss reduction and gradient-accumulation scheme.

NCCL enqueues its collective onto a CUDA stream. The host-side call can return before bytes have crossed the fabric. All ranks must participate in compatible collectives, in compatible order, with compatible tensor shapes. If one rank arrives late, the others wait. If one rank never arrives, the collective cannot complete without timeout or external failure handling.

5. Gradient accumulation changes the synchronization boundary

Large effective batches are often built from several microbatches. Each microbatch runs forward and backward, and its gradients accumulate into param.grad. The optimizer runs only after the final microbatch.

Without special handling, DDP would still synchronize every microbatch. Its no_sync() context suppresses those intermediate reductions so ranks communicate only on the accumulation boundary. This saves communication, but it also changes the failure and performance shape: a rank can execute several local backward passes before the global synchronization exposes skew between ranks.

Loss scaling must match the intended gradient average. Averaging each microbatch loss, summing gradients over accumulation steps, and averaging across ranks are three separate operations. Frameworks often hide the arithmetic, but a configuration error here changes the effective update even when the system runs perfectly.

6. The optimizer updates every replica locally

After the reducer has finished, every rank holds the same synchronized gradients. Each rank also started the step with the same parameters and optimizer state. Therefore every rank can run optimizer.step() locally and arrive at the same new parameters.

For Adam or AdamW, the update is more than subtracting the gradient. For each parameter, the optimizer reads the gradient and two running moment estimates, updates those estimates, computes a normalized update, and writes the parameter. Fused and multi-tensor optimizers combine many small updates into fewer GPU launches, but they still move substantial state through GPU memory.

The first optimizer step can create a memory surprise because many optimizers initialize state lazily. The model may fit through forward and backward, then fail when optimizer.step() allocates moment tensors for the first time. Later steps reuse them.

optimizer.zero_grad(set_to_none=True) then removes references to old gradient tensors rather than filling them with zeros. PyTorch can allocate fresh gradients as backward reaches each parameter, which can reduce memory operations and sometimes memory footprint. It does not erase optimizer state or model weights.

7. Checkpointing is outside the mathematical step but inside its wall time

A checkpoint is usually triggered after an optimizer boundary, when the weights and optimizer state describe a consistent step. The application serializes some combination of parameters, optimizer moments, learning-rate state, random-number-generator state, sampler position, and distributed metadata.

A synchronous checkpoint places that work on the critical path: later training cannot proceed until the required ranks finish staging and writing. An asynchronous checkpoint first copies stable state into separate buffers, then allows training to continue while a background path writes those buffers. This reduces exposed I/O time but consumes additional host memory and requires the application to control how many saves remain in flight. PyTorch's distributed-checkpoint documentation explicitly describes this memory tradeoff.

Checkpointing therefore belongs in a step-level performance model even though it does not happen every step. A 90-second save every 100 steps adds an average 900 milliseconds per step if none of it overlaps. The GPU timeline appears periodic: many normal steps, then a long cold interval with storage activity.

What each timing measurement really includes

The most useful view combines end-to-end step time with the lane timeline. A slow step is an outcome. A wide input gap, late bucket, long all-reduce tail, or oversized optimizer phase is an explanation.

A compact way to read a distributed step

When a step is slower than expected, ask these questions in dependency order:

1. Was the next batch ready? If not, the GPU waits before forward.
2. Did the host-to-device copy overlap? Pageable memory, stream placement, or NUMA distance can expose it.
3. Were forward and backward kernels continuously available? Gaps point toward the host, synchronization, or dynamic-shape overhead.
4. When did each gradient bucket become ready? A late bucket identifies the backward region holding up communication.
5. Did all ranks arrive together? Per-rank skew turns one local delay into a global delay.
6. How much communication remained after backward? That exposed tail, not total collective time, extends the critical path.
7. Did the optimizer or checkpoint create the peak? Both can add work and memory outside the forward/backward pair.

A distributed training step is complete only when every dependency required for the next update has completed—not when Python has finished enqueueing it.

Where RidgeScope fits

RidgeScope observes this execution below the training script: GPU activity and compute-pipe counters, CUDA kernel launches, NCCL collectives, rank behavior, job I/O, scheduler state, and the training progress printed by the framework. It does not need the model's source code to recognize the shape of the pipeline.

That distinction matters because most performance failures are empty spaces between valid operations. The model code may be correct. Every kernel may produce the right result. Yet a rank can wait for a batch, a bucket can arrive late, or a save can occupy the critical path. Understanding one step as an overlapped system makes those empty spaces legible.