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

NVIDIA Transformer Engine — uma biblioteca para acelerar o treinamento de modelos transformer em GPUs Ampere e mais recentes. Um tutorial publicado em 1º de agosto de 2026 no MarkTechPost mostra como conectar o TE a um modelo semelhante ao GPT, configurar a quantização FP8 e medir os ganhos de velocidade e memória.

O Que É o NVIDIA Transformer Engine

O NVIDIA Transformer Engine é uma biblioteca oficial para otimização de transformers em GPUs de nível Ampere ou superior. Ele fornece substituições drop-in para camadas padrão do PyTorch: `te.Linear`, `te.LayerNorm`, `te.LayerNormLinear`, `te.LayerNormMLP` e um bloco completo `te.TransformerLayer`. Todos os componentes funcionam com BF16 por padrão; os tensor cores FP8 estão disponíveis em placas com compute capability ≥ 8.9.

Componentes-chave da biblioteca:

• `te.Linear` — camada linear com suporte a tensor cores FP8
• `te.LayerNormLinear` — combina LayerNorm e projeção em um único kernel
• `te.LayerNormMLP` — bloco MLP fundido com normalização
• `te.TransformerLayer` — bloco transformer completo: atenção, FFN, normalização
• Fallback automático para PyTorch puro em GPUs sem suporte a TE

Como Funcionam o FP8 e o DelayedScaling

O modo FP8 é ativado por meio da receita `DelayedScaling` do módulo `transformer_engine.common.recipe`. A receita gerencia o histórico de amax dos tensores (`amax_history_len=16`), o algoritmo de computação (`amax_compute_algo="max"`) e o formato híbrido E4M3/E5M2: o primeiro formato garante precisão para as ativações, o segundo oferece maior alcance dinâmico para os gradientes.

A característica principal do delayed scaling é acumular estatísticas ao longo de 16 iterações antes de atualizar a escala do tensor. Isso estabiliza o treinamento nas etapas iniciais, quando a distribuição das ativações ainda não se estabeleceu.

"O

Transformer Engine combina kernels fundidos, escalonamento FP8 e caminhos de execução otimizados para Ampere, tornando-o uma escolha prática para o treinamento de modelos grandes", afirma a documentação oficial da NVIDIA.

No tutorial, os autores constroem um modelo compacto semelhante ao GPT, o `MiniGPT_TE`: quatro blocos `te.TransformerLayer`, dimensão oculta de 768, 12 cabeças de atenção, vocabulário de 96 tokens e comprimento de sequência de 256. O modelo é treinado em dados sintéticos determinísticos. As medições mostram que o modo FP8 em uma GPU de classe H100 reduz o consumo de pico de VRAM e o tempo por etapa em comparação com o BF16, enquanto a precisão permanece comparável — os autores verificam isso por meio de geração autorregressiva após o treinamento.

Requisitos de GPU

O Transformer Engine funciona apenas em GPUs com compute capability ≥ 8.0 (arquitetura Ampere). Os tensor cores FP8 estão disponíveis a partir de CC ≥ 8.9: Ada Lovelace (RTX 40xx) e Hopper (H100/H200). Em placas T4 e mais antigas, a biblioteca retorna ao PyTorch puro sem kernels fundidos.

• Mínimo para kernels TE: Ampere, CC ≥ 8.0 — A10, A100, RTX 30xx
• Mínimo para FP8: Ada Lovelace / Hopper, CC ≥ 8.9
• Instalação: `pip install transformer_engine[pytorch]`
• Dependências: CUDA 11.8+, PyTorch 2.0+

Segundo os autores do tutorial, uma GPU A100 ou L4 no Google Colab é suficiente para reprodução completa; na T4, apenas o modo de fallback sem kernels fundidos está disponível.

O Que Isso Significa

O NVIDIA Transformer Engine reduz a barreira de entrada para o treinamento de baixa precisão: em vez de implementar manualmente a quantização FP8, basta substituir as camadas padrão do PyTorch pelos equivalentes TE. Para equipes que treinam modelos grandes em clusters de GPU, esta é uma forma prática de reduzir o uso de memória e o tempo de iteração sem reformular a arquitetura.

Perguntas Frequentes

É Necessário um H100 para Usar o Transformer Engine?

Não. O NVIDIA Transformer Engine funciona em qualquer GPU com compute capability ≥ 8.0, incluindo A100 e RTX 3090. Os tensor cores FP8 estão disponíveis apenas com CC ≥ 8.9 — no H100 e no RTX 4090; nas demais placas, a biblioteca muda automaticamente para BF16.

Como Instalar o Transformer Engine?

Instale via pip: `pip install transformer_engine[pytorch]`. São necessários CUDA 11.8 ou superior e PyTorch 2.0+.