Loading…
How to Optimize Transformer-Based Models for Low-Precision Training
NVIDIA Developer BlogJonathan Mitchell
Summary
Accelerating transformer training requires understanding how low-precision formats such as FP8 and NVFP4 impact specific general matrix multiplication (GEMM) workloads. Transformer configurations do not explicitly reveal these shapes, making microbenchmarks necessary before committing to full training runs. NVIDIA Transformer Engine enables quantization and kernel dispatch across precisions, which can be evaluated in realistic autocast mode or kernel-only prequantized mode. Profiling the ESM2-15B model on NVIDIA B300 GPUs demonstrated that NVFP4 achieved a 1.79x blended forward propagation speedup over MXFP8 and up to 4.01x over BF16 in prequantized execution. Although large GEMM dimensions successfully overcome quantization overheads, dynamic scaling, Hadamard transforms, and kernel selection asymmetries moderate real-world gains.
Context
Large transformer models demand significant GPU hours and engineering time during training. While GPUs support low-precision formats such as FP8 and NVFP4, transformer configurations obscure exact runtime GEMM shapes, complicating precision selection and overhead estimation.
Approach / What changed
Using a microbenchmark tool with NVIDIA Transformer Engine, developers map transformer hyperparameters and batch sizes into exact M×K×N GEMM shapes for Fprop, Dgrad, and Wgrad stages, evaluating performance across precisions in autocast and prequantized modes on NVIDIA B300 GPUs.
Takeaways
- Prequantized microbenchmarking isolates raw kernel speedups from dynamic quantization overheads, showing NVFP4 versus BF16 jumping from 2.69x in autocast mode to 4.01x in kernel-only execution.
- FP8 DelayedScaling proved to be the fastest FP8 recipe on Blackwell at 23.76 ms per layer in autocast mode, outperforming FP8 CurrentScaling and MXFP8 by leveraging an amax-history approach.
- Quantized Dgrad passes run 4 to 5 percent slower than Fprop passes due to matrix aspect ratio changes that affect kernel selection, making the standard 2x Fprop timing approximation inaccurate.
Related reading
Make Long-Running NVIDIA TensorRT Engine Builds Observable and Cancelable in Python or C++
Long-running NVIDIA TensorRT engine builds can freeze workflows and waste GPU hours when they lack visibility or abort mechanisms. Developers can resolve this by subclassing the IProgressMonitor API in Python or C++ and attaching it to the builder configuration. The monitor tracks hierarchical build phases through start, step completion, and finish callbacks, requiring thread synchronization because TensorRT calls monitor methods across internal builder threads. Returning false from the step completion callback requests early build cancellation, which unwinds active phases in reverse order at the next step boundary. Beyond terminal rendering with ANSI escape codes, teams can stream these progress events to IDE protocols, HTTP service endpoints, or agent tool-call traces.
Michelle Horton