FP8 Training on AMD GPUs: Performance Boost with TorchTitan and TorchAO
13.4% throughput gain achieved through AMD optimizations for FP8 training
1Revolution in FP8 Training
Linear scaling demonstrated on AMD Instinct GPU clusters with over 1,000 GPUs. Thanks to Primus-Turbo library, TorchTitan now natively supports AMD GPUs. FP8 training delivers 13.4% higher throughput compared to BF16 on Llama3-8B model. This gain comes from faster FP8 matrix cores, not memory savings.
2Challenges in MoE Models
FP8 quantization initially introduced significant overhead in MoE architectures like DeepSeek-V3 671B. Fused Triton quantization kernels recovered 89% of this overhead. Individual kernel optimizations achieved up to 6.2x speedup per MoE layer. Grouped GEMM support simplified management of complex structures requiring per-expert scales.
3AMD-Specific Solutions
AMD's FNUZ (Finite, No NaN, Unsigned Zero) FP8 format differs from NVIDIA's e4m3fn. TorchAO initially caused silent errors by using incorrect maximum values. Hardware auto-detection was added and platform-specific loss baselines were established. Block-wise kernel support was implemented for MI300 and MI350 GPUs.
4Quantization Strategies
FP8 scaling can be applied at four levels: per-tensor (fastest), per-row (more accurate), per-block, or per-group packed with data (MXFP8). Each strategy was validated for AMD GPUs. Rowwise FP8 with high-precision weight-gradient recipe maintained similar memory usage to BF16 while improving throughput.
5Triton Kernel Optimization
The FP8 quantization pipeline initially required multiple kernel launches. Intermediate tensors were written to high-bandwidth memory between steps. Fusion combined absolute max computation, scale factor application, and FP8 conversion into a single kernel. This significantly reduced quantization overhead for MoE models.
6Grouped GEMM Support
In MoE models, each token is routed to a subset of experts. Grouped GEMM requires per-row and per-expert-column scales. FP8 grouped GEMM was enabled for AMD using Composable Kernel backend. A single kernel launch now handles quantization and expert routing.
7Accuracy-Performance Balance
AMD's e4m3fnuz format lacks NaN/Inf encodings, causing silent quality degradation from overflow. Hardware detection and platform-specific settings resolved this. Correct format selection became a correctness requirement. Accurate MI300X peak FLOPS reporting improved model efficiency measurements.
8Results and Impact
Rowwise FP8 training achieved 13.4% throughput gain over BF16 on Llama3-8B. 89% of quantization overhead was recovered on DeepSeek-V3 671B. Individual kernel optimizations reduced MoE layer processing from 7,290µs to 1,170µs. These advancements will critically reduce training costs for large language models.
9Industry Implications
FP8 support on AMD GPUs will intensify competition in high-performance computing. Open-source contributions are expanding the PyTorch ecosystem. MoE optimizations will guide future large-scale AI projects. These developments highlight the importance of hardware-software integration.