Transformer Engine: FP8 and FP4 Precision for Transformer Models on NVIDIA GPUs
A library for accelerating Transformer models on NVIDIA GPUs, including using 8-bit and 4-bit floating point (FP8 and FP4) precision on Hopper, Ada and Blackwell GPUs, to provide better performance with lower memory utilization in both training and inference.
At a glance
- What is it?
- NVIDIA's Transformer Engine is a Python and C++ library that wraps FP8, MXFP8 and NVFP4 kernels behind modules and an autocast context manager. It is a hardware-bound tool: Hopper, Ada and Blackwell only for the low-precision paths, and the framework-agnostic C++ API is the part that matters if you are not on PyTorch or JAX.
- Who is it for?
- Adopt Transformer Engine if you are training or serving Transformer models on Hopper, Ada or Blackwell hardware and want FP8 without writing scaling-factor management yourself. Do not adopt it if you are on Ampere or older and expecting the FP8 paths, or if you need a framework outside PyTorch and JAX and are unwilling to work against the C++ API.
- Can I use it commercially?
- Yes. Apache-2.0 is a permissive licence: you can use, modify and sell software built on it, as long as you keep its copyright and licence notices.
- Is it still maintained?
- Yes. The repository received new commits within the last day.
- What is it written in?
- Mainly Python, according to GitHub's language statistics.
Answers come from the project's GitHub data, last synced on September 15, 2026, and from our analysis. They are not legal advice.
DEEP OPEN-SOURCE ANALYSIS
The problem is not quantization, it is scaling-factor bookkeeping
Low-precision training fails in a specific way. You can cast tensors to FP8, but the dynamic range of activations and gradients shifts between layers and between steps, so a fixed scale either clips values or wastes mantissa bits. The README frames the library's purpose narrowly: modules internally maintain scaling factors and other values needed for FP8 training. That sentence is the whole pitch. Transformer Engine is not a model zoo or a training framework. It is a set of building blocks plus a recipe object that carries the scaling policy, so the caller does not have to track per-tensor amax history by hand. The audience is engineers already inside a training stack, typically working on models in the hundreds-of-billions-of-parameters range where the README says training and inference become memory and compute intensive. If your model fits comfortably in BF16 on one node, the library solves a problem you do not have.
How the autocast scope and recipe object fit together
The mechanism visible in the README is a context manager plus a recipe. You construct a module such as te.Linear, build a recipe with recipe.DelayedScaling(margin=0, fp8_format=recipe.Format.E4M3), then wrap the forward pass in te.autocast(enabled=True, recipe=fp8_recipe). The recipe is passed in at call time, not bound to the module at construction. That matters for anyone writing a training loop, because the scaling policy becomes an argument you can vary per step or per phase rather than a property baked into the layer. The README notes that all input args to the recipe are optional, so a default recipe exists. Format selection is explicit: E4M3 in the PyTorch example, HYBRID in the Flax example. Below the Python layer sits a framework-agnostic C++ library containing structs and kernels needed for FP8 support, which is what a non-PyTorch, non-JAX framework would integrate against. The README also states that optimizations such as fused kernels are provided for Transformer models, and that TE integrates with tensor, sequence and context parallelism plus MoE. Those integration claims are listed as highlights, not demonstrated in the README text.
Installation and the hardware gate you cannot configure around
The README links to an Installation section and a user guide at docs.nvidia.com rather than reproducing install steps inline, so the exact pip or source-build command is not in the material available here. What the README does state plainly is the hardware matrix, and it is the first thing to check. FP8 requires Hopper, Ada or Blackwell. MXFP8 and NVFP4 require Blackwell specifically. Optimizations across FP16 and BF16 are stated to work on Ampere and later. That is a hard boundary, not a tuning knob. A team on A100 hardware gets the fused kernels and the BF16 paths but not the FP8 support the library is named for. On the software side, the README shows two entry points: import transformer_engine.pytorch as te for PyTorch, and import transformer_engine.jax as te plus transformer_engine.jax.flax as te_flax for JAX with Flax. The Flax example constructs te_flax.DenseGeneral(features=HIDDEN) inside the autocast scope, which is a different placement from the PyTorch example, where the module is built before the scope opens. Whether that ordering is required or merely idiomatic in the example is not stated in the README, and it is worth confirming against the user guide before copying either pattern.
The C++ API is the escape hatch, and also the maintenance burden
The README describes a framework-agnostic C++ API that can be integrated with other deep learning libraries to enable FP8 support for Transformers. This is the part of the project with the least documentation in the supplied material. The Python surface is illustrated with runnable-looking snippets; the C++ surface is described in one sentence about structs and kernels. For a team on a framework NVIDIA has not integrated, that asymmetry is the real cost of adoption. You are not consuming a stable high-level interface, you are binding to kernel-level structs whose layout can move between releases. The release cadence visible in the metadata supports that concern: v2.17, v2.17.1 and v2.18 all landed within roughly a six-week window in mid-2026, and the README's news list runs through September 2026. A library shipping point releases at that rate is tracking new GPU architectures and new precision formats as they appear. That is the point of the project, but it means version pinning is not optional for anyone building on the C++ layer.
Where it is the wrong tool
Two cases stand out. The first is inference on non-NVIDIA accelerators. Nothing in the README suggests portability beyond NVIDIA GPUs, and the precision formats named are tied to specific NVIDIA architectures. If your deployment target is a different vendor, this library has nothing to offer you. The second is a team that wants FP8 as a post-training quantization step rather than a training-time precision. Transformer Engine's model is mixed-precision training and inference with scaling factors maintained during the forward and backward pass. The README's framing is training-first, with inference mentioned alongside it. If your workflow is take a finished BF16 checkpoint, quantize it offline, serve it, the recipe-and-autocast structure is more machinery than that task needs. A third, softer case: the README claims FP8 offers performance gains over FP16 with no degradation in accuracy. That claim is stated as a property of the format, not as a guarantee for your model. Accuracy behaviour under FP8 is model-dependent, and the README does not offer a procedure for detecting degradation. The Convergence section is linked in the header, which suggests the project has material on this, but the content is not in the supplied text.
What it is not competing with
The obvious comparison is with the automatic mixed precision already built into PyTorch, torch.cuda.amp. The difference in approach is the precision target and who manages the scale. AMP autocast selects between FP16 or BF16 and FP32 per operation, and GradScaler handles loss scaling for FP16. Transformer Engine targets FP8 and the newer 4-bit formats, and its modules carry the scaling factors internally rather than relying on a separate scaler object. That is a genuine architectural difference: in the TE example there is no GradScaler at all, just autocast with a recipe. The trade is that you give up framework-native modules for TE's replacements such as te.Linear, because the scaling bookkeeping lives inside those modules. You cannot get the FP8 path by wrapping an arbitrary nn.Module in autocast and hoping. That is why the README's examples all construct TE modules rather than stock PyTorch layers. If you want FP8 without swapping your layer implementations, this library does not offer that route.
Licence and upgrade cost
The project is Apache-2.0, which permits commercial use, modification and redistribution with the usual conditions around notices and patent grant. That is a permissive licence and it removes the licensing question for most commercial training stacks. It does not remove the upgrade question. Given the release cadence and the architecture-tracking nature of the work, expect to re-validate after upgrades rather than treating a version bump as routine. The specific things to re-check after any upgrade are the recipe constructor arguments (the README notes all are optional, which means defaults can change without breaking your call site) and the module constructor signatures for the TE layers you use. Pin the version in your environment file, and read the release notes linked in the README header before moving. The README's news list is a useful signal of what each release cycle is aimed at: entries through 2026 cover MoE fusion kernels, NVFP4 pretraining, FP8 reinforcement learning, and JAX with MaxText on Blackwell. If your workload is not in that set, a given release may not contain anything you need.
Editorial conclusion
Adopt Transformer Engine if you are training or serving Transformer models on Hopper, Ada or Blackwell hardware and want FP8 without writing scaling-factor management yourself. Do not adopt it if you are on Ampere or older and expecting the FP8 paths, or if you need a framework outside PyTorch and JAX and are unwilling to work against the C++ API. Before committing, verify three things: that your GPU generation matches the precision format you intend to use, that your installed PyTorch or JAX version is one the release notes list as supported for your TE version, and that your training loop tolerates the recipe object being threaded through autocast rather than configured globally.
Community notes