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

NVIDIA Transformer Engine — библиотека для ускорения обучения трансформерных моделей на GPU Ampere и новее. Туториал, опубликованный 1 августа 2026 года на MarkTechPost, показывает, как подключить TE к GPT-подобной модели, настроить FP8-квантование и измерить выигрыш по скорости и памяти.

Что такое NVIDIA Transformer Engine

NVIDIA Transformer Engine — официальная библиотека для оптимизации трансформеров на GPU уровня Ampere и выше. Она предоставляет drop-in замены стандартным слоям PyTorch: te.Linear, te.LayerNorm, te.LayerNormLinear, te.LayerNormMLP и полный блок te.TransformerLayer. Все компоненты работают с BF16 по умолчанию; на картах с compute capability ≥ 8.9 доступны FP8 tensor cores.

Ключевые компоненты библиотеки:

• te.Linear — линейный слой с поддержкой FP8 tensor cores
• te.LayerNormLinear — объединяет LayerNorm и проекцию в одном ядре
• te.LayerNormMLP — фьюзед MLP-блок с нормализацией
• te.TransformerLayer — полный блок трансформера: внимание, FFN, нормализация
• Автоматический откат на чистый PyTorch при GPU без поддержки TE

Как работает FP8 и DelayedScaling

FP8-режим активируется через рецепт DelayedScaling из transformer_engine.common.recipe. Рецепт управляет историей максимумов тензоров (amax_history_len=16), алгоритмом вычисления (amax_compute_algo="max") и гибридным форматом E4M3/E5M2: первый формат обеспечивает точность для активаций, второй — расширенный динамический диапазон для градиентов.

Ключевая особенность задержанного масштабирования — накопление статистики по 16 итерациям прежде, чем обновить масштаб тензора. Это стабилизирует обучение на ранних шагах, когда распределение активаций ещё не установилось.

«Transformer Engine сочетает фьюзед ядра, масштабирование FP8 и

Ampere-оптимизированные пути выполнения, что делает его практичным выбором для тренировки больших моделей», — говорится в официальной документации NVIDIA.

В туториале авторы строят компактную GPT-подобную модель MiniGPT_TE: четыре блока te.TransformerLayer, пространство размерностью 768, 12 голов внимания, словарь из 96 токенов, длина последовательности 256. Модель тренируется на детерминированных синтетических данных. Измерения показывают: FP8-режим на GPU класса H100 снижает пиковое потребление VRAM и время одного шага по сравнению с BF16, при этом точность остаётся сопоставимой — авторы проверяют её через авторегрессионную генерацию после тренировки.

Требования к GPU

Transformer Engine работает только на GPU с compute capability ≥ 8.0 (архитектура Ampere). FP8 tensor cores доступны начиная с CC ≥ 8.9: это Ada Lovelace (RTX 40xx) и Hopper (H100/H200). На картах T4 и более старых библиотека откатывается на чистый PyTorch без фьюзед-ядер.

• Минимум для TE-ядер: Ampere, CC ≥ 8.0 — A10, A100, RTX 30xx
• Минимум для FP8: Ada Lovelace / Hopper, CC ≥ 8.9
• Установка: pip install transformer_engine[pytorch]
• Зависимости: CUDA 11.8+, PyTorch 2.0+

По данным авторов туториала, для полного воспроизведения достаточно GPU A100 или L4 в Google Colab; на T4 доступен только откатной режим без фьюзед-ядер.

Что это значит

NVIDIA Transformer Engine снижает порог входа в низкоточное обучение: вместо ручной реализации FP8-квантования достаточно заменить стандартные PyTorch-слои на TE-аналоги. Для команд, тренирующих большие модели на GPU-кластерах, это практический способ снизить расход памяти и время итерации без переработки архитектуры.

Частые вопросы

Нужен ли H100, чтобы использовать Transformer Engine?

Нет. NVIDIA Transformer Engine работает на любом GPU с compute capability ≥ 8.0, включая A100 и RTX 3090. FP8 tensor cores доступны только с CC ≥ 8.9 — на H100 и RTX 4090; на остальных картах библиотека автоматически переключается на BF16.

Как установить Transformer Engine?

Установка через pip: pip install transformer_engine[pytorch]. Требуются CUDA 11.8 или выше и PyTorch 2.0+.