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:
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.
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);
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).
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.
The Ring algorithm works in two phases over N
GPUs connected in a logical ring:
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
).
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.
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.
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.
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.
Y = X * W
where W
is split column-wise. This requires an All-Reduce after every linear layer.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.
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.
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);
}
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! 🚀