transformer-engine-cu12

What it is and what it does

Transformer Engine is a library that accelerates Transformer model training and inference on NVIDIA GPUs by providing optimized kernels and low-precision arithmetic support. It enables 8-bit floating-point (FP8) training on Hopper, Ada, and Ampere GPUs, and adds support for MXFP8 and NVFP4 formats on Blackwell GPUs. The library handles scaling factors and precision management internally, allowing developers to use a simple autocast API similar to mixed-precision training frameworks.

The package integrates with popular frameworks through framework-specific modules and provides a C++ API for integration with other deep learning libraries. It includes fused operations, support for distributed training patterns (tensor/sequence/context parallelism), and Mixture-of-Experts (MoE) optimizations. Installation requires a compatible NVIDIA GPU, CUDA 12.1+, cuDNN 9.3+, and a C++ compiler with C++17 support; wheels are pre-compiled for specific architectures (x86_64, aarch64) on manylinux_2_28.

Use it for:

• Train large language models with FP8 precision to reduce memory footprint and increase throughput without accuracy loss
• Accelerate Mixture-of-Experts (MoE) model training using fused kernels and low-precision formats
• Run inference on Transformer models with reduced latency and memory using optimized GPU kernels
• Implement mixed-precision training workflows with automatic scaling factor management
• Deploy multimodal or biology-focused Transformer models on Blackwell GPUs using NVFP4 for maximum efficiency

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

Accelerates Transformer model training and inference on NVIDIA GPUs using low-precision formats (FP8, MXFP8, NVFP4) with optimized kernels and mixed-precision APIs.

Yes, if you are training or serving Transformer models on supported NVIDIA GPUs (Ampere or newer) and have the required CUDA/cuDNN stack. The library is actively maintained, has no known vulnerabilities, and offers significant performance and memory benefits through low-precision training. However, verify license terms before commercial use and ensure your system meets the strict hardware and software prerequisites (CUDA 12.1+, cuDNN 9.3+, C++17 compiler).

Install

transformer-engine-cu12 on PyPI

pip install transformer-engine-cu12

uv add transformer-engine-cu12

poetry add transformer-engine-cu12

Installing transformer-engine-cu12

Before you install

Medium install friction due to compiled wheels for specific GPU architectures (aarch64, x86_64) and CUDA 12 requirement. Active maintenance with recent release (3 days old). Requires Python 3.10+, CUDA 12.1+ (or 12.8+ for Blackwell), cuDNN 9.3+, and GCC 9+ or Clang 10+ with C++17 support.

License in practice

License treatment is unclear—no SPDX identifier or raw license text provided in metadata. Verify licensing terms before use, particularly for commercial deployment.

Quickstart

pip install transformer-engine-cu12

import transformer_engine.pytorch as te
from transformer_engine.common import recipe

model = te.Linear(768, 3072, bias=True)
fp8_recipe = recipe.DelayedScaling(margin=0, fp8_format=recipe.Format.E4M3)

with te.autocast(enabled=True, recipe=fp8_recipe):
    out = model(inp)

Requires NVIDIA GPU (Hopper, Ada, Ampere, or Blackwell), CUDA 12.1+, cuDNN 9.3+, and Linux or WSL2. Compiled wheels are architecture-specific; verify your system matches manylinux_2_28 with x86_64 or aarch64.

Verify before relying

• Exact performance gains (speedup percentages, memory savings) on different GPU architectures and model sizes
• Accuracy degradation (if any) when using FP8, MXFP8, or NVFP4 formats versus standard precision
• Compatibility with specific framework versions beyond the stated Python 3.10+ requirement
• Support status and maintenance timeline for older GPU architectures (Ampere) versus newer ones (Blackwell)

Package facts

Evidence: transformer_engine_cu12-2.18.0-py3-none-manylinux_2_28_aarch64.whl; transformer_engine_cu12-2.18.0-py3-none-manylinux_2_28_x86_64.whl

Tags

More Artificial Intelligence packages

Further reading