TransformerEngine

Why we included this project

Transformer Engine is NVIDIA's own library for making low-precision training practical on its GPUs. It ships drop-in PyTorch and JAX modules for building Transformer layers with FP8 support, and on Blackwell it adds MXFP8 and NVFP4, so you can cut memory use and speed up both training and inference without hand-writing the scaling-factor bookkeeping that mixed precision normally demands. NVIDIA runs many of its own LLM and MoE training jobs on it, and it plugs into NeMo and MaxText, so the optimizations have clearly been exercised at scale. On Hopper, Ada, and Blackwell GPUs the fused kernels and parallelism-aware features (MoE, tensor/sequence/context parallelism) are directly useful, and the framework-agnostic C++ API lets other deep learning libraries bolt on FP8 support. It is not a turnkey training tool, so expect to work within your existing PyTorch or JAX stack rather than adopting a new workflow.