AMD and Meta Push TorchAO FP8 Training to 6.2x Faster on Instinct GPUs

FP8 training -- where matrix multiplications run in 8-bit floating point instead of the usual 16-bit -- has been a reliable way to squeeze more throughput out of NVIDIA hardware for a while. AMD's Instinct GPUs support FP8 too, but getting it to actually work well in PyTorch required a non-trivial amount of plumbing. That plumbing is now done. AMD and Meta engineers have upstreamed a full suite of FP8 optimizations into mainline TorchAO and TorchTitan, so AMD Instinct GPUs get competitive FP8 performance out of the box with no AMD-specific installs required.

The hidden correctness bug that had to come first

Before any performance work could land, there was a silent correctness problem to fix. AMD Instinct GPUs use a variant of FP8 called FNUZ (Finite, No NaN, Unsigned Zero), specifically e4m3fnuz, which has a maximum representable value of 240. NVIDIA's FP8 format, e4m3fn, has a max of 448. TorchAO was hardcoded to NVIDIA's format, so on AMD hardware it was computing scaling factors against the wrong ceiling.

Because e4m3fnuz has no NaN/Inf encodings, the overflow did not raise an error -- it degraded model quality instead. Selecting the correct format is therefore a correctness requirement, not a tuning option. The fix was hardware auto-detection, so TorchAO now selects the correct format automatically. Alongside this, the team fixed MFU (Model FLOP Utilization) reporting to use the correct MI300X peak FLOPS, and added platform-specific loss baselines for FNUZ numerics.

Dense models: a clean 13.4% win

Rowwise FP8 with a high-precision weight-gradient recipe -- where the weight-update GEMM stays in BF16 while the forward and gradient-input GEMMs use FP8 -- delivers a 13.4% throughput gain over BF16, with peak memory nearly identical (~39 GB). The win comes from faster FP8 matrix cores, not memory savings. This was measured on Llama3-8B across 8×MI300X GPUs with FSDP2 and torch.compile.

The "rowwise" part matters here. Rowwise scaling is better at handling outliers than tensorwise scaling, so these recipes are different points on the accuracy vs performance curve. Tensorwise is fastest but coarsest; rowwise trades a small amount of compute for better numerical fidelity. For most training workloads, rowwise is the right default.

MoE models: where FP8 gets complicated

Dense models are relatively straightforward -- every linear layer has the same shape, so a single quantization strategy applies uniformly. Mixture-of-Experts (MoE) models like DeepSeek-V3 and Llama 4 are different. They route each input token to a subset of "expert" sub-networks, which produces variable-size batches that must be processed through a grouped GEMM -- a single kernel call that handles multiple matrix multiplications of different sizes simultaneously. FP8 quantization for grouped GEMM requires per-row scales on activations, per-expert-column scales on weights, and an offset tensor routing rows to the correct expert. That's a lot more moving parts than dense training.

When AMD first enabled FP8 on MoE shapes, the quantization overhead was severe enough to erase the FP8 speedup entirely. The root cause was structural: the FP8 quantization pipeline in TorchAO converts tensors through a multi-step chain:

1. Compute per-row/column absolute max (absmax)
2. Derive the scale factor and apply it
3. Clamp and cast to FP8

Each step was a separate GPU kernel launch, and each step materialized an intermediate tensor to HBM (High Bandwidth Memory) before the next step could start. For MoE models with dozens of expert weight tensors per layer, these extra memory round-trips dominated the FP8 overhead. The math was cheap; the data movement was not.

Three levels of kernel surgery

The team attacked the data movement problem at three levels of granularity, each compounding on the last.

Level 1 -- Launch fewer kernels. The backward pass had two compounding problems: a .t().contiguous().t() pattern forced a full tensor copy through HBM just to convert weight layout for GEMM compatibility, and the multi-step scale-and-cast chain launched separate kernels with intermediate tensors materialized between them. Removing the redundant transpose and fusing the chain into single Triton kernels delivered a 4.2x backward pass throughput improvement on 8×MI300X GPUs with DeepSeek-MoE-16B shapes. On the forward path, quantizing expert weights launched five generic kernels per call -- with 24 calls per step, that added ~90 ms/step of overhead. A single fused Triton kernel replaced the entire chain, collapsing five launches into one and letting surrounding GEMMs issue sooner. On DeepSeek-V3 671B shapes, the forward-pass kernel fusion alone recovered 89% of the quantization overhead (5,996 → 7,027 tok/s vs 7,156 BF16 baseline on 8×MI325X GPU).

Level 2 -- Make each kernel move memory efficiently. The colwise scales kernel used in the backward pass had non-coalesced memory writes: consecutive SIMD lanes wrote to addresses far apart in memory, each triggering a separate memory transaction. The fix was to transpose the output tile through LDS (Local Data Share, AMD's on-chip scratchpad) before storing, and add a fused single-pass variant that eliminates a redundant HBM read. Result: 6.2x speedup per MoE layer on MI300X with DeepSeek-V3 671B shapes (7,290μs → 1,170μs).

Level 3 -- Strip synchronization the hardware never needed. Triton's atomic operations default to acquire-release memory ordering, which inserts memory fences before and after every atomic. On AMD GPUs, these fences are expensive and unnecessary for commutative reductions like absmax. Switching to relaxed ordering -- guarded by a torch.version.hip check so NVIDIA behavior is unchanged -- removed a class of hidden synchronization overhead.

What didn't work

The team also tried expanding the Triton autotune search space for MoE FP8 kernels from 1 to 8-16 candidate configurations, expecting the 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 takeaway: autotuning search spaces should be shaped by hardware constraints (wavefront size, LDS capacity, register pressure), not expanded to more candidates by default.

The numbers, side by side

Workload	Optimization	Result
Llama3-8B (dense)	Rowwise FP8 vs BF16	+13.4% throughput
DeepSeek-MoE-16B	Backward transpose removal + fusion	4.2x backward pass
DeepSeek-V3 671B	Colwise scales coalescing	6.2x per MoE layer (7,290→1,170μs)
DeepSeek-V3 671B	Forward pass fusion	+17% end-to-end; recovers 89% of FP8 gap

Available now, no AMD-specific setup required

All contributions have been merged into mainline pytorch/ao and pytorch/torchtitan. TorchAO's FP8 training support is integrated into both TorchTitan and TorchTune, so all scaling recipes can be used for pre-training and fine-tuning with minimal setup. If you have AMD Instinct MI300X, MI325X, or MI350X GPUs, upgrading TorchAO and TorchTitan is all you need -- no separate AMD library, no custom fork.

The supported scaling strategies are:

• Tensorwise -- one scale per tensor, fastest, coarsest
• Rowwise -- one scale per row, better accuracy, recommended default
• Blockwise -- one scale per fixed-size tile, now supported on MI300 and MI350
• MXFP8 -- microscaling format, grouped alongside the data

Work continues on next-generation hardware. The team is developing MXFP8 grouped GEMM and quantization kernels for forward and backward passes on MI355X GPUs. The broader implication is significant: the assumption that FP8 training is an NVIDIA-only story no longer holds. AMD Instinct hardware now participates in the same upstream PyTorch training stack, with competitive numbers and a clear optimization roadmap.