Life of a Training Step

concept

One training step, followed through all nine layers. This is the fastest way to see why the layers have to be designed together: a single gradient touches every one of them, and a weak hop anywhere slows down every GPU in the job.

The setup

Two HGX nodes, 8 GPUs each, 16 GPUs in total. The model is trained with FSDP (fully sharded data parallel) across all 16 GPUs. Each GPU has its own 400G RDMA NIC on a rail-optimized fabric. The control plane has already placed the pods and attached the GPUs and NICs (see from pod spec to data path); from here on it plays no part.

The step, hop by hop

sequenceDiagram
    autonumber
    participant S as Storage (S)
    participant APP as Training loop (L1)
    participant FW as PyTorch FSDP (L2/L3)
    participant K as Kernels (L4)
    participant G as GPU + HBM (L9)
    participant N as NCCL (L5)
    participant NV as NVLink (L8)
    participant T as Verbs + GPUDirect (L6)
    participant F as Fabric (L7)
    participant P as Peer node GPUs

    S->>APP: Read next batch (data loader)
    APP->>FW: forward(batch)
    FW->>N: All-gather this layer's weight shards
    N->>NV: Intra-node part over NVLink
    N->>T: Inter-node part as RDMA writes
    T->>F: NIC reads GPU memory directly (GPUDirect RDMA)
    F->>P: Packets cross the fabric to the peer NIC
    FW->>K: Launch matmul and attention kernels
    K->>G: Compute on tensor cores, read and write HBM
    APP->>FW: loss.backward()
    FW->>K: Gradient kernels
    FW->>N: Reduce-scatter gradients
    N->>NV: Intra-node reduce over NVLink
    N->>T: Inter-node reduce over RDMA
    T->>F: GPU memory to fabric, no CPU copy
    F->>P: Peers receive and sum their shard
    APP->>FW: optimizer.step()
    FW->>G: Update the local shard of weights
    APP->>S: Every N steps, write a checkpoint

Steps 3–7 and 12–16 repeat for every layer of the model, and FSDP overlaps them with compute: while layer k is computing, the weights for layer k+1 are already being gathered. When the overlap works, communication is hidden. When it doesn't, GPUs wait.

What each layer contributes

Layer	Role in this step	If it is weak
S Storage	Feeds the next batch; absorbs checkpoints	GPUs idle at the start of each step, or for minutes at every checkpoint
L1 Application	The training loop: batch, forward, backward, step	Small batches leave too little compute to hide communication behind
L2/L3 Framework and parallelism	FSDP decides what is gathered and reduced, and when	Wrong sharding or no prefetch turns overlapped communication into blocking waits
L4 Kernels	Turn each operation into GPU work	Unfused or low-precision-unaware kernels waste tensor-core time
L5 NCCL	Runs the all-gather and reduce-scatter collectives	A poor algorithm or wrong NIC choice cuts bandwidth several-fold
L6 Transport	RDMA verbs with GPUDirect: NIC reads GPU memory directly	Without GPUDirect, data bounces through host memory; without RDMA, NCCL falls back to TCP
L7 Fabric	Carries inter-node traffic	Congestion, pause storms or one bad link slow the whole ring
L8 Interconnect	NVLink for intra-node traffic; PCIe from GPU to NIC	A GPU paired with a NIC on the other CPU socket crosses the inter-socket link on every byte
L9 GPU and HBM	Does the math, holds weights, gradients and activations	Out of memory forces smaller batches or more sharding, which means more communication

The lesson

Every GPU waits for the slowest participant in each collective. A single mis-paired NIC on one node, a single flapping link, or a single slow storage mount sets the speed of all 16 GPUs. That is why this site treats the stack as one system and why co-design is the method rather than tuning each layer alone.

When a step is slower than expected, the slow collectives decision tree walks the same path from the bottom up.