AMD Upstreams FP8 Training Optimizations for TorchTitan and TorchAO

AMD and Meta/PyTorch engineers have merged FP8 training optimizations for AMD Instinct GPUs into the upstream TorchTitan and TorchAO repositories. The changes include native support for AMD's FP8 format, grouped GEMM for Mixture-of-Experts models on ROCm, and a Triton kernel fusion pipeline that reduces quantization overhead. Teams with AMD Instinct GPUs can now get these gains by upgrading TorchAO and TorchTitan, with no AMD-specific installation required.

FP8 format correctness on AMD hardware

AMD Instinct GPUs implement a variant of FP8 called e4m3fnuz, which has a max value of 240 and no NaN or Inf encodings. TorchAO initially did not support this format, so it computed scales against a different max value, producing silently wrong results: tensors were scaled into a range that exceeded the hardware's representable values, clipping activations and corrupting gradients. Because e4m3fnuz has no NaN or Inf encodings, the overflow did not raise an error; it degraded model quality instead. The fix was hardware auto-detection so TorchAO selects the correct FP8 dtype and max value automatically. Additional correctness fixes include reporting correct MI300X peak FLOPS for accurate MFU numbers and adding platform-specific loss baselines for FNUZ numerics.

MoE grouped GEMM and kernel fusion

Mixture-of-Experts models route each token to a subset of experts, producing variable-size batches that require grouped GEMM with per-row scales on activations, per-expert-column scales on weights, and an offset tensor for routing. The team enabled FP8 grouped GEMM on ROCm by adapting the quantization pipeline to use the correct dtype and dispatch via the Composable Kernel backend.

The FP8 quantization pipeline in TorchAO originally converted tensors through a multi-step chain: compute per-row or column absolute max, derive the scale factor, apply it, then clamp and cast to FP8. Each step was a separate kernel launch, materializing intermediate tensors to HBM between steps. For MoE models with dozens of expert weight tensors per layer, these round-trips dominated the FP8 overhead. The optimizations reduce data movement at three levels: launching fewer kernels, making each kernel move memory efficiently, and removing unnecessary synchronization.

Level 1 fused the multi-step chain into single Triton kernels. On the backward pass, removing redundant transpose copies and fusing the scale-and-cast chain delivered a 4.2x backward pass throughput improvement on DeepSeek-MoE-16B with 8xMI300X GPUs. On the forward pass, quantizing expert weights launched five generic kernels per call, and with 24 calls per step, this added about 90 ms per step of overhead. Replacing the entire chain with a single fused Triton kernel that parallelizes across both experts and output-dimension blocks delivered a 17% end-to-end throughput improvement on DeepSeek-V3 671B with 8xMI325X GPUs, recovering 89% of the BF16 to FP8 gap.

Level 2 fixed non-coalesced memory writes in the colwise scales kernel used in the backward pass. Consecutive SIMD lanes wrote to addresses K bytes apart, each triggering a separate memory transaction. Transposing the output tile through LDS before storing, plus a fused single-pass variant, delivered a 6.2x speedup per MoE layer on MI300X with DeepSeek-V3 671B shapes.

Level 3 stripped synchronization that the hardware never needed. Triton's atomic operations default to acquire-release memory ordering, which on AMD GPUs inserts memory fences before and after every atomic, expensive for commutative reductions. Switching to relaxed ordering on AMD GPUs, guarded by a torch.version.hip check, left NVIDIA behavior unchanged.

What didn't work: autotuning

The team expanded the Triton autotune search space for MoE FP8 kernels from 1 to 8 to 16 candidate configurations, expecting wider search to find faster tile sizes on AMD's wavefront-based architecture. Benchmarking on Llama 4 shapes on MI300X showed no measurable improvement, and the extra configs increased first-iteration compile time. The change was reverted. The takeaway: autotuning search spaces should be shaped by hardware constraints like wavefront size, LDS capacity, and register pressure, not expanded to more candidates by default.

Work continues on next-generation hardware, including MXFP8 grouped GEMM and quantization kernels for forward and backward passes on MI355X GPUs. The kernel fusion pipeline continues with further Triton optimizations.