FP8 training without the scaling-factor bookkeeping

What it does TransformerEngine is NVIDIA’s library for training and running Transformers in low-precision number formats: FP8 on Hopper, Ada, and Blackwell GPUs, plus MXFP8 and NVFP4 on Blackwell. It ships optimized building blocks for common Transformer layers, exposed through an autocast-style Python API for both PyTorch and JAX/Flax, plus a framework-agnostic C++ API for wiring FP8 into other deep learning libraries. On the PyTorch side it also supports FlashAttention-2 and FlashAttention-3, preferring the latter when both are present.

The interesting bit FP8’s hard part was never the arithmetic — it’s that every tensor needs scaling factors tracked through forward and backward passes, or the numbers overflow and underflow. TE’s modules maintain those internally, so low-precision training feels like ordinary mixed precision rather than a numerics research project. The pedigree is real: NVIDIA says Nemotron 3 was trained with NVFP4 on it, and DeepL, LMSYS, and the Ling 2.0 team have all published accounts of using it.

Key highlights

• FP8 on Hopper, Ada, and Blackwell; MXFP8 and NVFP4 on Blackwell; FP16/BF16 optimizations on Ampere and newer
• Autocast-like API for PyTorch and JAX/Flax, with scaling factors managed inside the modules
• Framework-agnostic C++ API for adding FP8 support to other DL libraries
• Plays well with MoE, tensor/sequence/context parallelism, and fused kernels
• FlashAttention-2 and FlashAttention-3 support in PyTorch

Caveats

• FP8 requires Compute Capability 8.9+ — Ampere owners get FP16/BF16 optimizations only
• Source builds have documented sharp edges: FlashAttention-2 compilation can exhaust RAM, and a C++ ABI mismatch with PyTorch built outside NVIDIA’s containers causes undefined-symbol import errors
• Linux is the supported platform (WSL2 gets limited support), and the conda JAX package is still listed as “coming soon”

Verdict If you train or serve large Transformers on Ada, Hopper, or Blackwell hardware, this is the mainstream route to FP8/FP4 — the same machinery behind NVIDIA’s own training runs. On older GPUs or non-NVIDIA silicon, the interesting parts simply don’t apply.