Transformer Engine (TE) is a Python library by NVIDIA for accelerating Transformer models on NVIDIA GPUs. It enables lower precision training and inference, notably supporting 8-bit (FP8) and 4-bit (NVFP4) floating point precision on Hopper, Ada, and Blackwell GPUs, leading to better performance and reduced memory utilization. It provides highly optimized building blocks for popular Transformer architectures and an automatic mixed precision-like API for PyTorch and JAX. The current version is 2.13.0, with an active release cadence, often aligning with new NVIDIA hardware and software advancements.
pip install --no-build-isolation transformer-engine-cu12[pytorch]Verified import paths — ran on the pinned version, not inferred.
This quickstart demonstrates how to use `transformer_engine.pytorch.Linear` with FP8 autocasting. Ensure you have PyTorch and a compatible CUDA environment set up. The `fp8_autocast` context manager automatically handles FP8 quantization for supported operations within its scope.
Review the Transformer Engine 2.2 release notes for `InferenceParams` and `DelayedScaling` API updates. Adjust code to use new required arguments and method calls.
Update C++ code to use the non-packed fused attention APIs. Consult Transformer Engine's C++ API documentation for the correct alternatives.
Always include `--no-build-isolation` in your `pip install` commands for Transformer Engine to ensure future compatibility and prevent potential build issues. E.g., `pip install --no-build-isolation transformer-engine-cu12[pytorch]`.
Ensure PyTorch and Transformer Engine are built with the same C++ ABI. Rebuilding PyTorch from source with a matching ABI might be necessary, or use NVIDIA NGC Docker containers where these dependencies are pre-configured.
Ensure your CUDA environment is at least 12.8 or higher when installing from PyPI. If you must use an older CUDA 12.x version (e.g., 12.1), consider installing Transformer Engine from source and explicitly managing CUDA paths during the build process, or use an NGC container.