Welcome back to the final (for now) chapter of our GPU optimization saga! In Part 1, we mastered single-GPU matrix multiplication and built the transformer block. In Part 2, we implemented backpropagation, the AdamW optimizer, and mixed-precision training.
But here is the harsh reality: a single GPU—even an H100 or an MI300X—can barely hold a 70B parameter model in VRAM, let alone train it within a human lifetime. To train state-of-the-art LLMs, we must span hundreds or thousands of GPUs.
In this third part, we will tear down the walls of a single device and venture into the world of distributed training. We will implement All-Reduce for data parallelism, explore Ring-AllReduce algorithms, dive into ZeRO (Zero Redundancy Optimizer) sharding, and write real HIP/C++ code using NCCL (NVIDIA) and RCCL (AMD) to synchronize gradients across nodes.
By the end of this part, you will understand:
- How collective communication (All-Reduce) works under the hood.
- How to shard optimizer states and gradients (ZeRO Stage 1 & 2).
- How to write a fully distributed training loop that scales linearly.
- The difference between Data Parallelism, Model Parallelism, and Pipeline Parallelism.
Prerequisites: Completion of Parts 1 & 2, a cluster with at least 2-4 GPUs (they don't have to be same vendor, but we will write portable HIP code), and nccl/rccl libraries installed.
1. The Communication Backbone – NCCL vs. RCCL
CUDA uses NCCL (NVIDIA Collective Communications Library), while ROCm uses RCCL (ROCm Collective Communications Library). Fortunately, they share the exact same API signatures with a prefix change (nccl vs rccl). We can unify them using preprocessor macros.
// Unified header selection
#ifdef __HIP_PLATFORM_AMD__
#include <rccl/rccl.h>
#define COMM_ID rcclUniqueId
#define COMM_INIT rcclCommInitRank
#define COMM_ALL_REDUCE rcclAllReduce
#define COMM_GET_ERROR rcclGetErrorString
#else
#include <nccl.h>
#define COMM_ID ncclUniqueId
#define COMM_INIT ncclCommInitRank
#define COMM_ALL_REDUCE ncclAllReduce
#define COMM_GET_ERROR ncclGetErrorString
#endif
Initializing the communicator requires a unique ID that is broadcasted to all ranks (typically via MPI, but we can use simple environment variables for a single node):
COMM_ID id;
if (rank == 0) { // rank 0 generates the ID
COMM_GET_UNIQUE_ID(&id);
}
// In a real cluster, you broadcast this via MPI_Bcast.
// For single-node, we just pass it directly.
COMM_COMM_T comm;
COMM_INIT(&comm, world_size, id, rank);
2. Data Parallelism – The All-Reduce Operation
In Data Parallelism (DP) , each GPU holds a full copy of the model. We feed different micro-batches to each GPU, compute local gradients (dW_local), and then average them across all GPUs.
The mathematical operation is:
dW_global = (1 / world_size) * Σ dW_local_i
This is exactly an All-Reduce operation with the SUM operator (we just divide by world_size separately or use AVG if supported).
2.1 Naïve All-Reduce (Parameter Server)
A naive approach uses a central "server" GPU to receive all gradients, average them, and send them back. This creates a bottleneck. Instead, we use the Ring-AllReduce algorithm, which is bandwidth-optimal.
2.2 Implementing Ring-AllReduce in your Kernel
The Ring algorithm works in two phases over N GPUs connected in a logical ring:
- Scatter-Reduce (N-1 steps): Each GPU sends a unique chunk of its gradient to the next GPU while receiving another chunk, accumulating the sum.
- All-Gather (N-1 steps): GPUs forward the now-aggregated chunks around the ring until every GPU has the full sum.
Instead of writing this from scratch, we use the highly optimized COMM_ALL_REDUCE:
void sync_gradients(float* d_gradient, int num_elements,
ncclComm_t comm, hipStream_t stream) {
// In-place All-Reduce: sums gradients across all GPUs
// Result is stored in d_gradient on every rank.
COMM_ALL_REDUCE((const void*)d_gradient, // sendbuff
(void*)d_gradient, // recvbuff (in-place)
num_elements,
COMM_FLOAT, // data type
COMM_SUM, // operation
comm, stream);
// Average the sum to get the mean gradient
int world_size;
COMM_COMM_COUNT(comm, &world_size);
float inv_world = 1.0f / world_size;
scale_kernel<<<(num_elements+255)/256, 256, 0, stream>>>(d_gradient, inv_world, num_elements);
}
Note on AMD: RCCL uses the exact same function signatures. Just replace nccl with rccl in your linker flags (-lrccl vs -lnccl).
3. ZeRO (Zero Redundancy Optimizer) – Sharding the Memory
Data Parallelism is great, but each GPU still stores the entire model, optimizer states, and gradients. For 175B parameters, this requires ~3.2TB just for weights (in FP32). ZeRO, introduced by Microsoft, shards these components across GPUs to reduce memory footprint.
3.1 ZeRO Stage 1 – Sharding the Optimizer States
In Part 2, we stored m (momentum) and v (variance) for AdamW. These are the same size as the model weights. In Stage 1, we partition these optimizer states: GPU 0 holds the optimizer states for parameters 0 to N/2, and GPU 1 holds the rest.
Implementation modification to our AdamW kernel:
Instead of launching the kernel over all parameters, we launch it only over the local partition.
__global__ void adamw_update_sharded_kernel(float* W, float* dW, float* m, float* v,
int total_params, int rank, int world_size,
...) {
int idx = blockIdx.x * blockDim.x + threadIdx.x + rank * (total_params / world_size);
if (idx >= total_params) return;
// ... (same update logic as Part 2)
}
During the backward pass, we must All-Gather the updated weights so that every GPU has the full updated model before the next forward pass. This adds a communication step after the optimizer.
3.2 ZeRO Stage 2 – Sharding Gradients
Stage 2 goes further: it shards the gradients (dW) as well. Before the All-Reduce, we only reduce gradients that belong to the local GPU's partition. This reduces the communication volume by up to half.
Pseudocode for the training step with ZeRO-2:
// Forward: All GPUs have full weights (via All-Gather after previous step)
forward_pass();
// Backward: Compute local gradients (full size)
backward_pass();
// Reduce-Scatter: Each GPU sums only its assigned partition of the gradients
reduce_scatter_gradients(dW, local_partition, ...);
// Update: Update only local weights and local optimizer states
adamw_update_local(W_local, dW_local, m_local, v_local, ...);
// All-Gather: Sync the updated weights so everyone has the full model
all_gather_weights(W_full, W_local, ...);
This is the core of modern frameworks like DeepSpeed and FairScale.
4. Hybrid Parallelism – Model + Data Parallelism
For truly massive models (e.g., > 1 Trillion parameters), even ZeRO-3 isn't enough. We need Model Parallelism (MP) , where the linear layers themselves are split across GPUs.
-
Tensor Parallelism (TP): Splitting the weight matrices row-wise or column-wise. For example,
Y = X * WwhereWis split column-wise. This requires an All-Reduce after every linear layer. -
Pipeline Parallelism (PP): Splitting the layers across GPUs (e.g., GPU 0 holds layers 1-10, GPU 1 holds layers 11-20). This requires sending activations (forward) and gradients (backward) between GPUs via
P2P(Peer-to-Peer) communication.
Here is how we launch a P2P send/receive in NCCL/RCCL:
// GPU 0 sends activations to GPU 1
if (rank == 0) {
COMM_SEND(d_activation, size, comm, 1, stream);
} else if (rank == 1) {
COMM_RECV(d_activation, size, comm, 0, stream);
}
The 3D Parallelism Paradigm (DP + TP + PP) is the secret sauce behind training GPT-4 and Llama 3. You apply TP within a node (using NVLink/Infinity Fabric), PP across nodes, and DP across groups of nodes.
5. Optimizing the Communication Overlap
Communication is expensive. The best way to hide it is Overlap. While the GPU is computing the backward pass for layer 10, we can start the All-Reduce for the gradients of layer 1 in the background.
We achieve this using Streams and Event-based Synchronization:
hipStream_t compute_stream, comm_stream;
hipStreamCreate(&compute_stream);
hipStreamCreate(&comm_stream);
// During backward pass:
for (int layer = 0; layer < num_layers; ++layer) {
// Compute gradient for this layer on compute_stream
backward_layer<<<..., compute_stream>>>(...);
// Record an event when gradient is ready
hipEventRecord(grad_ready_event, compute_stream);
// Make comm_stream wait for the event
hipStreamWaitEvent(comm_stream, grad_ready_event, 0);
// Launch All-Reduce on comm_stream (async)
COMM_ALL_REDUCE(..., comm_stream);
}
This way, the GPU spends almost zero idle time waiting for network packets; the computation fills the gaps.
6. The Fully Distributed Training Loop (Final Code)
Combining everything we've learned—ZeRO-2 sharding, overlapping communication, and mixed-precision—here is the skeleton of a single training step for a distributed environment:
void distributed_train_step(float* d_model_weights, ... , int rank, int world_size,
ncclComm_t comm, hipStream_t comp_stream, hipStream_t comm_stream) {
// 1. Forward Pass (Full model on each GPU)
forward_pass(d_model_weights, d_activations, ..., comp_stream);
// 2. Backward Pass (Compute full gradients)
backward_pass(d_model_weights, d_gradients, ..., comp_stream);
// 3. Reduce-Scatter Gradients (ZeRO-2)
hipEventRecord(grad_ready, comp_stream);
hipStreamWaitEvent(comm_stream, grad_ready, 0);
int local_size = total_params / world_size;
COMM_REDUCE_SCATTER(d_gradients, d_gradients_local, ... , comm_stream);
hipEventRecord(comm_done, comm_stream);
hipStreamWaitEvent(comp_stream, comm_done, 0);
// 4. Update local weights & optimizer states (AdamW)
adamw_update_local<<<..., comp_stream>>>(d_local_weights, d_gradients_local,
d_momentum_local, d_variance_local, ...);
// 5. All-Gather to sync the full model
hipEventRecord(update_done, comp_stream);
hipStreamWaitEvent(comm_stream, update_done, 0);
COMM_ALL_GATHER(d_local_weights, d_full_weights, ... , comm_stream);
// 6. Synchronize the main stream at the very end
hipStreamSynchronize(comp_stream);
}
Conclusion
We have just crossed the finish line of the GPU optimization marathon. Starting from a simple vector addition in Part 1, we built a single-GPU transformer, added a full training stack with mixed precision in Part 2, and finally shattered the single-GPU ceiling by implementing distributed data parallelism, ZeRO sharding, and overlapping communication using NCCL/RCCL in this final part.
You now possess the architectural blueprints that power every major AI company's training cluster. Whether you are running on a 4-GPU workstation at home or a 4,000-GPU supercomputer, the principles remain exactly the same: reduce communication, maximize compute, and shard everything that can be sharded.
What lies beyond?
If there is demand for a Part 4, we could explore FP8 Quantization for even faster training, Fused Multi-Head Attention (Flash Attention) kernels to optimize memory bandwidth further, or Distributed Checkpointing to save sharded models without crashing the filesystem.
Let me know in the comments below what excites you the most!
Thank you for sticking with me through this deep dive. If you ran into any errors while implementing these kernels, or if you have a specific scenario (e.g., using AMD MI300X with RCCL vs NVIDIA H100 with NCCL) that you want me to elaborate on, drop a comment—I personally respond to all of them.
Until next time, keep your warps converged and your bandwidth saturated!
See ya on the next adventure! 🚀
Top comments (0)