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+.