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.