NVIDIA Transformer Engine: как ускорить обучение трансформеров в 2-3 раза с FP8 и BF16

NVIDIA Transformer Engine (TE) — библиотека, которая объединяет fused GPU-ядра, вычисления в BF16 и аппаратное FP8 для ускорения трансформеров. В туториале показан полный пайплайн: от установки и детекции GPU (определяется compute capability, поддержка TE и FP8) до сборки и обучения компактной GPT-подобной модели на 768-мерных эмбеддингах с 4 слоями.

Ключевые компоненты TE: te.Linear, te.LayerNorm, te.LayerNormLinear, te.LayerNormMLP и te.TransformerLayer. Для FP8 настраивается delayed-scaling рецепт с гибридным форматом E4M3/E5M2, историей amax и алгоритмом max. Если GPU не поддерживает TE (например, T4), автоматически включается fallback на чистый PyTorch.

Собранная модель MiniGPT_TE тренируется на синтетических последовательностях. Сравниваются режимы BF16 и FP8 по времени выполнения и пиковому потреблению GPU-памяти. После обучения модель валидируется через автогрессивную генерацию. Результаты показывают, что FP8 даёт ускорение в 2-3 раза без потери качества на синтетике.