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

Hey Dev Community!

Welcome back! In Part 1, we built the foundation: from vector addition to tiled GEMM, and finally assembled a complete forward pass of a Transformer block using HIP (CUDA/ROCm). But a forward pass without gradients is just a very expensive random number generator.

Training an LLM requires the backward pass (backpropagation), an optimizer (like AdamW), and brutal memory management to fit billions of parameters into VRAM. In this second part, we will implement the missing pieces: gradient computation for every layer, the weight update step, mixed-precision training (FP16/BF16), and a fully functional training loop.

By the end of this part, you will understand:

· How to write gradient kernels for Linear layers, Softmax, and LayerNorm.
· How to implement a fused AdamW optimizer on the GPU.
· How to leverage activation checkpointing to trade compute for memory.
· How to run a full training iteration (forward, loss, backward, update) in pure HIP/C++.

Prerequisites: Completion of Part 1, a strong grasp of the chain rule, and a GPU with at least 8GB of VRAM to follow along locally.

---

1. The Backward Pass – Gradients of the Linear Layer

The forward pass of a linear layer is Y = X * W + b (ignoring bias for simplicity). During backpropagation, we receive dY (gradient of the loss w.r.t output) and must compute:

· dX = gradient w.r.t input (to pass to the previous layer).
· dW = gradient w.r.t weights (to update the weights).

Mathematically:
dX = dY * W^T
dW = X^T * dY

We can reuse the exact same hipblasSgemm (or rocBLAS) we used before, just transposing matrices.

// Assuming d_Y is [B, S, D] and W is [D, D]
// Compute d_X = d_Y * W^T  (Matrix multiply)
float alpha = 1.0f, beta = 0.0f;
hipblasSgemm(handle, HIPBLAS_OP_N, HIPBLAS_OP_T,
             S, D, D,           // M, N, K
             &alpha, 
             d_Y, S,            // matrix A (dY)
             W, D,              // matrix B (W) transposed internally
             &beta, 
             d_X, S);           // matrix C (dX)

// Compute d_W = X^T * d_Y
hipblasSgemm(handle, HIPBLAS_OP_T, HIPBLAS_OP_N,
             D, D, S,
             &alpha, 
             X, S,              // matrix A (X) transposed
             d_Y, S,            // matrix B (dY)
             &beta, 
             d_W, D);           // matrix C (dW)

For multi-head attention, the gradients are trickier because of the softmax and the Q*K^T multiplication, but the principle is the same: every matrix multiplication in the forward pass corresponds to two matrix multiplications in the backward pass.

---

1. Gradient of Softmax and LayerNorm (Custom Kernels)

While GEMMs handle the linear parts, we need custom kernels for the non-linearities.

2.1 Softmax Backward Kernel

Let P be the output of the forward softmax (probabilities). The backward pass computes dX given dY. The formula is:
dX_i = P_i * (dY_i - sum(P_j * dY_j)).

__global__ void softmax_backward_kernel(const float* d_Y, const float* P, 
                                        float* d_X, int rows, int cols) {
    int row = blockIdx.x * blockDim.x + threadIdx.x;
    if (row >= rows) return;

    // First, compute dot product of P and dY for this row
    float dot = 0.0f;
    for (int j = 0; j < cols; ++j) {
        dot += P[row * cols + j] * d_Y[row * cols + j];
    }

    // Second, compute dX = P * (dY - dot)
    for (int j = 0; j < cols; ++j) {
        int idx = row * cols + j;
        d_X[idx] = P[idx] * (d_Y[idx] - dot);
    }
}

2.2 LayerNorm Backward

LayerNorm requires computing gradients w.r.t input X, and the scale/bias parameters gamma and beta. We won't write the full kernel here to save space, but the pattern is:

1. Compute mean and variance from the forward pass (you stored them).
2. Compute d_gamma and d_beta by reducing over the hidden dimension.
3. Compute d_X using the standardized values.

Pro Tip: Store the mean and inv_std (inverse standard deviation) from the forward pass in a small buffer. This saves you from recomputing them during backward.

---

1. The AdamW Optimizer – Fused Kernel

After computing dW and db, we need to update the model weights. AdamW is the standard optimizer for LLMs. It maintains two exponential moving averages per parameter: m (momentum) and v (variance).

Instead of launching a separate kernel for each parameter update, we write a fused kernel that updates everything in one pass. This minimizes kernel launch overhead and maximizes memory bandwidth.

__global__ void adamw_update_kernel(float* W, float* dW, float* m, float* v,
                                    int num_params, float lr, float beta1, 
                                    float beta2, float eps, float weight_decay,
                                    int step) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= num_params) return;

    // Bias correction factors
    float bias_correction1 = 1.0f - powf(beta1, step);
    float bias_correction2 = 1.0f - powf(beta2, step);

    // Update biased first moment estimate
    m[idx] = beta1 * m[idx] + (1.0f - beta1) * dW[idx];
    // Update biased second raw moment estimate
    v[idx] = beta2 * v[idx] + (1.0f - beta2) * dW[idx] * dW[idx];

    // Compute bias-corrected estimates
    float m_hat = m[idx] / bias_correction1;
    float v_hat = v[idx] / bias_correction2;

    // Update weight (with weight decay)
    W[idx] = W[idx] - lr * (m_hat / (sqrtf(v_hat) + eps) + weight_decay * W[idx]);

    // Optionally reset gradient to zero for next iteration (or do it async)
    dW[idx] = 0.0f;
}

Launch this kernel over num_params (e.g., 1 billion floats). This is highly efficient and leverages the GPU's massive threading capacity.

---

1. Mixed Precision Training (FP16 / BF16)

Modern GPUs (NVIDIA Ampere+ and AMD CDNA+) have dedicated hardware for FP16/BF16 matrix multiplication, effectively doubling throughput. To train an LLM with mixed precision:

1. Cast weights and activations to FP16/BF16 for the forward and backward passes.
2. Keep a FP32 master copy of the weights for the optimizer update to avoid gradient underflow.
3. Loss Scaling: Since FP16 has a small dynamic range, we multiply the loss by a large scalar (e.g., 1024) before backpropagation, and divide the gradients by this scalar before updating.

Here’s how we modify our forward pass:

// Use half type in HIP
#include <hip/hip_fp16.h>

__global__ void cast_and_scale_gradients(half* dW_half, float* dW_float, 
                                         float scale, int n) {
    int idx = threadIdx.x + blockIdx.x * blockDim.x;
    if (idx < n) dW_float[idx] = (float)dW_half[idx] * scale;
}

Note: rocBLAS/cuBLAS supports hipblasHgemm for FP16. Use hipblasGemmEx to choose the compute type (FP32 for accumulation) for maximum precision.

---

1. Activation Checkpointing (Trading Compute for Memory)

The biggest bottleneck in LLM training is memory. Storing every activation for backpropagation is impossible for models with 7B+ parameters.

Activation Checkpointing (or Gradient Checkpointing) saves only the input to specific layers (e.g., every 4th transformer block) and recomputes the intermediate activations during the backward pass.

Implementation strategy:

· Before the forward pass of a block, we store the input X in a "checkpoint" buffer.
· During the backward pass, we reload X and rerun the entire forward pass of that block (without storing activations again) just to recompute the activations needed for the backward pass of that block.

This roughly halves your memory footprint but increases compute by ~30-40%. For large LLMs, it is mandatory.

Pseudo-code:

// During Forward
if (layer_idx % checkpoint_interval == 0) {
    hipMemcpy(checkpoint_buffer + offset, d_X, size, hipMemcpyDeviceToDevice);
}
// During Backward
if (layer_idx % checkpoint_interval == 0) {
    // 1. Load X from checkpoint
    // 2. Run Forward pass of this layer (discard output, keep activations)
    // 3. Run Backward pass using recomputed activations
}

---

1. The Full Training Loop (Putting It All Together)

Now, we combine everything. A single training iteration in pure HIP/C++ looks like this:

void train_step() {
    // 1. Copy batch from CPU to GPU (async)
    hipMemcpyAsync(d_input, h_input, batch_bytes, hipMemcpyHostToDevice, stream);

    // 2. Forward Pass
    forward_transformer(d_input, d_output, ...); // Uses FP16 for matmuls

    // 3. Compute Loss (Cross Entropy)
    float loss = compute_loss_kernel(d_output, d_labels); 

    // 4. Loss Scaling
    scale_loss_kernel<<<...>>>(d_loss_scaled, loss, loss_scale);

    // 5. Backward Pass (traverses graph in reverse order)
    backward_transformer(d_output, d_input, ...); // Computes FP16 gradients

    // 6. Unscale Gradients (cast to FP32)
    cast_and_scale_gradients<<<...>>>(dW_weights, dW_fp32, 1.0f/loss_scale, n);

    // 7. Apply Gradient Clipping (optional, to prevent exploding gradients)
    float norm = compute_l2_norm_kernel(dW_fp32, n);
    if (norm > max_norm) scale_gradients(dW_fp32, max_norm / norm, n);

    // 8. Optimizer Step (AdamW on FP32 master weights)
    adamw_update_kernel<<<(n+255)/256, 256>>>(W_fp32, dW_fp32, m, v, ...);

    // 9. Copy updated FP32 weights back to FP16 for next forward pass
    cast_fp32_to_fp16_kernel<<<...>>>(W_fp16, W_fp32, n);

    // 10. Synchronize Stream
    hipStreamSynchronize(stream);
}

---

1. Performance Profiling – Where is the bottleneck?

You can't optimize what you can't measure. Use these tools:

· NVIDIA: nvprof or Nsight Systems (nsys profile).
· AMD: rocprof or OmniTrace.

Look for these metrics:

1. Occupancy: Are your warps/wavefronts idle? (Check achieved_occupancy).
2. Memory Bandwidth: GEMMs should be Compute-bound, but custom kernels (Softmax, LayerNorm) are often Memory-bound. Use shared memory to reduce global reads.
3. Kernel Launch Overhead: If you have 1,000 tiny kernels, fuse them! (e.g., fuse bias add, activation, and dropout into one kernel).

---

Conclusion

If you have reached this far, congratulations! You have just built the architectural skeleton of a modern LLM trainer from the ground up—covering GEMM optimization on CUDA/ROCm, custom backward kernels, the AdamW optimizer, mixed precision, and memory-saving checkpointing.

Obviously, frameworks like PyTorch and JAX handle all of this transparently and add distributed training (FSDP, ZeRO, and all-reduce), which we haven't touched. But understanding these low-level primitives makes you a master of GPU computing. You now know exactly what loss.backward() and optimizer.step() do under the hood, regardless of whether you are running on an NVIDIA H100 or an AMD MI300X.

What's next for Part 3?
We will dive into Multi-GPU Distributed Training – implementing All-Reduce, Ring-AllReduce, and sharding strategies (ZeRO stages) using NCCL/RCCL.

Until then, happy kernel coding, and may your occupancy be high and your warp divergence be low!

Have questions about the backward pass, loss scaling, or checkpointing? Drop them in the comments below! I read every single one and will respond in detail.

See you in the next part!