Advanced GPU Optimization: How can I tech an LLM with CUDA and ROCm? - Part 3

A developer's third installment in a GPU optimization series details implementing distributed training for large language models using CUDA and ROCm, covering All-Reduce, Ring-AllReduce, and ZeRO sharding with NCCL and RCCL libraries. The post provides portable HIP/C++ code to synchronize gradients across GPUs, enabling training of models too large for a single device.

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 🚀