NVIDIA Transformer Engine: фьюзед-ядра, FP8 и BF16 для обучения трансформеров

NVIDIA Transformer Engine — a library for accelerating transformer model training on Ampere and newer GPUs. A tutorial published on August 1, 2026 on MarkTechPost shows how to connect TE to a GPT-like model, configure FP8 quantization, and measure speed and memory gains.

What Is NVIDIA Transformer Engine

NVIDIA Transformer Engine is an official library for optimizing transformers on Ampere-class and higher GPUs. It provides drop-in replacements for standard PyTorch layers: `te.Linear`, `te.LayerNorm`, `te.LayerNormLinear`, `te.LayerNormMLP`, and a full `te.TransformerLayer` block. All components work with BF16 by default; FP8 tensor cores are available on cards with compute capability ≥ 8.9.

Key library components:

• `te.Linear` — linear layer with FP8 tensor core support
• `te.LayerNormLinear` — combines LayerNorm and projection in a single kernel
• `te.LayerNormMLP` — fused MLP block with normalization
• `te.TransformerLayer` — full transformer block: attention, FFN, normalization
• Automatic fallback to plain PyTorch on GPUs without TE support

How FP8 and DelayedScaling Work

FP8 mode is activated via the `DelayedScaling` recipe from `transformer_engine.common.recipe`. The recipe manages tensor amax history (`amax_history_len=16`), the computation algorithm (`amax_compute_algo="max"`), and the hybrid E4M3/E5M2 format: the first format provides precision for activations, the second provides extended dynamic range for gradients.

The key feature of delayed scaling is accumulating statistics over 16 iterations before updating the tensor scale. This stabilizes training in early steps when the activation distribution has not yet settled.

"Transformer

Engine combines fused kernels, FP8 scaling, and Ampere-optimized execution paths, making it a practical choice for training large models," states the official NVIDIA documentation.

In the tutorial, the authors build a compact GPT-like model `MiniGPT_TE`: four `te.TransformerLayer` blocks, a hidden dimension of 768, 12 attention heads, a vocabulary of 96 tokens, and a sequence length of 256. The model is trained on deterministic synthetic data. Measurements show that FP8 mode on an H100-class GPU reduces peak VRAM consumption and per-step time compared to BF16, while accuracy remains comparable — the authors verify this through autoregressive generation after training.

GPU Requirements

Transformer Engine only works on GPUs with compute capability ≥ 8.0 (Ampere architecture). FP8 tensor cores are available starting from CC ≥ 8.9: these are Ada Lovelace (RTX 40xx) and Hopper (H100/H200). On T4 cards and older, the library falls back to plain PyTorch without fused kernels.

• Minimum for TE kernels: Ampere, CC ≥ 8.0 — A10, A100, RTX 30xx
• Minimum for FP8: Ada Lovelace / Hopper, CC ≥ 8.9
• Installation: `pip install transformer_engine[pytorch]`
• Dependencies: CUDA 11.8+, PyTorch 2.0+

According to the tutorial authors, an A100 or L4 GPU in Google Colab is sufficient for full reproduction; on T4 only the fallback mode without fused kernels is available.

What This Means

NVIDIA Transformer Engine lowers the entry barrier for low-precision training: instead of manually implementing FP8 quantization, it is enough to replace standard PyTorch layers with TE equivalents. For teams training large models on GPU clusters, this is a practical way to reduce memory usage and iteration time without reworking the architecture.

Frequently Asked Questions

Do You Need an H100 to Use Transformer Engine?

No. NVIDIA Transformer Engine works on any GPU with compute capability ≥ 8.0, including the A100 and RTX 3090. FP8 tensor cores are only available with CC ≥ 8.9 — on the H100 and RTX 4090; on other cards the library automatically switches to BF16.

How to Install Transformer Engine?

Install via pip: `pip install transformer_engine[pytorch]`. CUDA 11.8 or higher and PyTorch 2.0+ are required.