{"slug": "advanced-gpu-optimization-how-can-i-tech-an-llm-with-cuda-and-rocm-part-3", "title": "Advanced GPU Optimization: How can I tech an LLM with CUDA and ROCm? - Part 3", "summary": "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.", "body_md": "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.\n\nBut 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**.\n\nIn 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.\n\nBy the end of this part, you will understand:\n\n**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`\n\n/`rccl`\n\nlibraries installed.\n\nCUDA 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`\n\nvs `rccl`\n\n). We can unify them using preprocessor macros.\n\n```\n// Unified header selection\n#ifdef __HIP_PLATFORM_AMD__\n    #include <rccl/rccl.h>\n    #define COMM_ID rcclUniqueId\n    #define COMM_INIT rcclCommInitRank\n    #define COMM_ALL_REDUCE rcclAllReduce\n    #define COMM_GET_ERROR rcclGetErrorString\n#else\n    #include <nccl.h>\n    #define COMM_ID ncclUniqueId\n    #define COMM_INIT ncclCommInitRank\n    #define COMM_ALL_REDUCE ncclAllReduce\n    #define COMM_GET_ERROR ncclGetErrorString\n#endif\n```\n\nInitializing 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):\n\n```\nCOMM_ID id;\nif (rank == 0) { // rank 0 generates the ID\n    COMM_GET_UNIQUE_ID(&id);\n}\n// In a real cluster, you broadcast this via MPI_Bcast.\n// For single-node, we just pass it directly.\nCOMM_COMM_T comm;\nCOMM_INIT(&comm, world_size, id, rank);\n```\n\nIn **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`\n\n), and then average them across all GPUs.\n\nThe mathematical operation is:\n\n`dW_global = (1 / world_size) * Σ dW_local_i`\n\nThis is exactly an **All-Reduce** operation with the `SUM`\n\noperator (we just divide by world_size separately or use `AVG`\n\nif supported).\n\nA 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.\n\nThe Ring algorithm works in two phases over `N`\n\nGPUs connected in a logical ring:\n\nInstead of writing this from scratch, we use the highly optimized `COMM_ALL_REDUCE`\n\n:\n\n```\nvoid sync_gradients(float* d_gradient, int num_elements, \n                    ncclComm_t comm, hipStream_t stream) {\n    // In-place All-Reduce: sums gradients across all GPUs\n    // Result is stored in d_gradient on every rank.\n    COMM_ALL_REDUCE((const void*)d_gradient,  // sendbuff\n                    (void*)d_gradient,         // recvbuff (in-place)\n                    num_elements,\n                    COMM_FLOAT,                // data type\n                    COMM_SUM,                  // operation\n                    comm, stream);\n\n    // Average the sum to get the mean gradient\n    int world_size;\n    COMM_COMM_COUNT(comm, &world_size);\n    float inv_world = 1.0f / world_size;\n    scale_kernel<<<(num_elements+255)/256, 256, 0, stream>>>(d_gradient, inv_world, num_elements);\n}\n```\n\n*Note on AMD:* RCCL uses the exact same function signatures. Just replace `nccl`\n\nwith `rccl`\n\nin your linker flags (`-lrccl`\n\nvs `-lnccl`\n\n).\n\nData 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.\n\nIn Part 2, we stored `m`\n\n(momentum) and `v`\n\n(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`\n\nto `N/2`\n\n, and **GPU 1** holds the rest.\n\n**Implementation modification to our AdamW kernel:**\n\nInstead of launching the kernel over all parameters, we launch it only over the local partition.\n\n```\n__global__ void adamw_update_sharded_kernel(float* W, float* dW, float* m, float* v,\n                                            int total_params, int rank, int world_size,\n                                            ...) {\n    int idx = blockIdx.x * blockDim.x + threadIdx.x + rank * (total_params / world_size);\n    if (idx >= total_params) return;\n    // ... (same update logic as Part 2)\n}\n```\n\nDuring 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.\n\nStage 2 goes further: it shards the gradients (`dW`\n\n) 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.\n\n*Pseudocode for the training step with ZeRO-2:*\n\n```\n// Forward: All GPUs have full weights (via All-Gather after previous step)\nforward_pass();\n\n// Backward: Compute local gradients (full size)\nbackward_pass();\n\n// Reduce-Scatter: Each GPU sums only its assigned partition of the gradients\nreduce_scatter_gradients(dW, local_partition, ...);\n\n// Update: Update only local weights and local optimizer states\nadamw_update_local(W_local, dW_local, m_local, v_local, ...);\n\n// All-Gather: Sync the updated weights so everyone has the full model\nall_gather_weights(W_full, W_local, ...);\n```\n\nThis is the core of modern frameworks like DeepSpeed and FairScale.\n\nFor 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.\n\n`Y = X * W`\n\nwhere `W`\n\nis split column-wise. This requires an All-Reduce after every linear layer.`P2P`\n\n(Peer-to-Peer) communication.Here is how we launch a P2P send/receive in NCCL/RCCL:\n\n```\n// GPU 0 sends activations to GPU 1\nif (rank == 0) {\n    COMM_SEND(d_activation, size, comm, 1, stream);\n} else if (rank == 1) {\n    COMM_RECV(d_activation, size, comm, 0, stream);\n}\n```\n\n**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.\n\nCommunication 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.\n\nWe achieve this using **Streams** and **Event-based Synchronization**:\n\n```\nhipStream_t compute_stream, comm_stream;\nhipStreamCreate(&compute_stream);\nhipStreamCreate(&comm_stream);\n\n// During backward pass:\nfor (int layer = 0; layer < num_layers; ++layer) {\n    // Compute gradient for this layer on compute_stream\n    backward_layer<<<..., compute_stream>>>(...);\n\n    // Record an event when gradient is ready\n    hipEventRecord(grad_ready_event, compute_stream);\n    // Make comm_stream wait for the event\n    hipStreamWaitEvent(comm_stream, grad_ready_event, 0);\n\n    // Launch All-Reduce on comm_stream (async)\n    COMM_ALL_REDUCE(..., comm_stream);\n}\n```\n\nThis way, the GPU spends almost zero idle time waiting for network packets; the computation fills the gaps.\n\nCombining 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:\n\n```\nvoid distributed_train_step(float* d_model_weights, ... , int rank, int world_size, \n                            ncclComm_t comm, hipStream_t comp_stream, hipStream_t comm_stream) {\n\n    // 1. Forward Pass (Full model on each GPU)\n    forward_pass(d_model_weights, d_activations, ..., comp_stream);\n\n    // 2. Backward Pass (Compute full gradients)\n    backward_pass(d_model_weights, d_gradients, ..., comp_stream);\n\n    // 3. Reduce-Scatter Gradients (ZeRO-2)\n    hipEventRecord(grad_ready, comp_stream);\n    hipStreamWaitEvent(comm_stream, grad_ready, 0);\n\n    int local_size = total_params / world_size;\n    COMM_REDUCE_SCATTER(d_gradients, d_gradients_local, ... , comm_stream);\n    hipEventRecord(comm_done, comm_stream);\n    hipStreamWaitEvent(comp_stream, comm_done, 0);\n\n    // 4. Update local weights & optimizer states (AdamW)\n    adamw_update_local<<<..., comp_stream>>>(d_local_weights, d_gradients_local, \n                                              d_momentum_local, d_variance_local, ...);\n\n    // 5. All-Gather to sync the full model\n    hipEventRecord(update_done, comp_stream);\n    hipStreamWaitEvent(comm_stream, update_done, 0);\n    COMM_ALL_GATHER(d_local_weights, d_full_weights, ... , comm_stream);\n\n    // 6. Synchronize the main stream at the very end\n    hipStreamSynchronize(comp_stream);\n}\n```\n\nWe 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.\n\nYou 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.\n\n**What lies beyond?**\n\nIf 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.\n\nLet me know in the comments below what excites you the most!\n\nThank 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.\n\nUntil next time, keep your warps converged and your bandwidth saturated!\n\nSee ya on the next adventure! 🚀", "url": "https://wpnews.pro/news/advanced-gpu-optimization-how-can-i-tech-an-llm-with-cuda-and-rocm-part-3", "canonical_source": "https://dev.to/javadinteger/advanced-gpu-optimization-how-can-i-tech-an-llm-with-cuda-and-rocm-part-3-39d8", "published_at": "2026-08-10 23:47:54+00:00", "updated_at": "2026-08-11 00:15:28.841827+00:00", "lang": "en", "topics": ["machine-learning", "large-language-models", "ai-infrastructure", "developer-tools"], "entities": ["NVIDIA", "AMD", "NCCL", "RCCL", "CUDA", "ROCm", "HIP", "ZeRO"], "alternates": {"html": "https://wpnews.pro/news/advanced-gpu-optimization-how-can-i-tech-an-llm-with-cuda-and-rocm-part-3", "markdown": "https://wpnews.pro/news/advanced-gpu-optimization-how-can-i-tech-an-llm-with-cuda-and-rocm-part-3.md", "text": "https://wpnews.pro/news/advanced-gpu-optimization-how-can-i-tech-an-llm-with-cuda-and-rocm-part-3.txt", "jsonld": "https://wpnews.pro/news/advanced-gpu-optimization-how-can-i-tech-an-llm-with-cuda-and-rocm-part-3.jsonld"}}