FP8 Training on AMD GPUs with TorchTitan and TorchAO: Upstreaming Performance Improvements – PyTorch
AMD and Meta/PyTorch engineers upstreamed FP8 training optimizations for AMD Instinct GPUs directly into TorchAO and TorchTitan, giving competitive FP8 performance out of the box without any AMD-specific installs. On dense models, rowwise FP8 delivers a 13.4% throughput gain over BF16 on Llama3-8B. For MoE architectures like DeepSeek-V3 671B, FP8 quantization initially added heavy overhead due to memory-bound kernel launches and HBM round-trips; fused Triton kernels recovered 89% of that overhead, with the colwise scales fix alone delivering a 6.2x speedup per MoE layer. Key work included adding hardware auto-detection for AMD's e4m3fnuz FP8 format (fixing silent numerical corruption from mismatched max values), enabling FP8 grouped GEMM on ROCm via Composable Kernel, and a three-level Triton fusion pipeline reducing kernel launches, improving memory coalescing, and relaxing unnecessary atomic memory fences on AMD hardware. An autotuning search-space expansion was tried and reverted after showing no benefit. Work continues on MXFP8 grouped GEMM for MI355X GPUs.
Table of contents
AMD FP8 format in TorchAOScaling FP8 to MoE ArchitecturesTriton Kernel OptimizationSummary and Next StepsAdditional ResourcesQuestions this post answers
Does TorchTitan support FP8 training on AMD Instinct GPUs now?
Yes, AMD's FP8 optimizations from Primus-Turbo have been upstreamed directly into pytorch/ao and pytorch/torchtitan, giving competitive FP8 performance out of the box with nothing AMD-specific to install. Teams get the gains by upgrading TorchAO and TorchTitan. Rowwise FP8 delivers a 13.4% throughput gain over BF16 on Llama3-8B on 8xMI300X GPUs. daily.dev surfaces upstream framework changes like this for teams tuning GPU training throughput.
Why does FP8 training on AMD GPUs give silently wrong results if the wrong number format is used?
AMD Instinct GPUs use the e4m3fnuz FP8 format, which has a max representable value of 240 and no NaN/Inf encodings. If a library computes scales against a different max value (as TorchAO initially did, assuming NVIDIA's e4m3fn), tensors get scaled beyond the hardware's representable range, clipping activations and corrupting gradients silently rather than raising an error. TorchAO now auto-detects the platform and selects the correct FP8 dtype automatically. developers debugging GPU numerics issues follow fixes like this one on daily.dev.
How much do fused Triton kernels reduce FP8 quantization overhead for MoE models like DeepSeek-V3?
Fused Triton kernels recovered 89% of the FP8 quantization overhead on DeepSeek-V3 671B MoE shapes. Replacing a five-kernel eager chain with a single fused kernel gave a 17% end-to-end throughput improvement (5,996 to 7,027 tok/s) on 8xMI325X GPUs, and fixing non-coalesced memory writes in the colwise scales kernel produced a 6.2x speedup per MoE layer on MI300X. engineers optimizing MoE training performance track kernel-level wins like these on daily.dev.