NVIDIA Transformer Engine: kernels fusionados, FP8 e BF16 para treinar transformers

O 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 integrar o TE a um modelo estilo GPT, configurar a quantização FP8 e medir o ganho de velocidade e de memória.

O que é o NVIDIA Transformer Engine

O NVIDIA Transformer Engine é a biblioteca oficial para otimizar transformers em GPUs da classe Ampere ou superior. Ela oferece substitutos drop-in para as camadas padrão do PyTorch: `te.Linear`, `te.LayerNorm`, `te.LayerNormLinear`, `te.LayerNormMLP` e o bloco completo `te.TransformerLayer`. Todos os componentes usam BF16 por padrão; em placas com compute capability ≥ 8.9, os tensor cores FP8 ficam disponíveis.

Principais componentes 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 fusionado com normalização
• `te.TransformerLayer` — bloco transformer completo: atenção, FFN e normalização
• Retorno automático ao PyTorch puro em GPUs sem suporte ao TE

Como funcionam o FP8 e o DelayedScaling

O modo FP8 é ativado pela receita `DelayedScaling`, do módulo `transformer_engine.common.recipe`. A receita controla o histórico de máximos dos tensores (`amax_history_len=16`), o algoritmo de cálculo (`amax_compute_algo="max"`) e o formato híbrido E4M3/E5M2: o primeiro garante precisão para as ativações, e o segundo oferece maior faixa dinâmica para os gradientes.

A principal característica do escalonamento atrasado é acumular estatísticas ao longo de 16 iterações antes de atualizar a escala do tensor. Isso estabiliza o treinamento nos primeiros passos, quando a distribuição das ativações ainda não se firmou.

«O

Transformer Engine combina kernels fusionados, escalonamento FP8 e caminhos de execução otimizados para Ampere, o que o torna uma escolha prática para treinar modelos grandes», diz a documentação oficial da NVIDIA.

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

Requisitos de GPU

O Transformer Engine só funciona 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 recorre ao PyTorch puro, sem kernels fusionados.

• Mínimo para os kernels do 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 basta para reproduzir tudo; na T4, só está disponível o modo de retorno ao PyTorch, sem kernels fusionados.

O que isso significa

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

Perguntas frequentes

É preciso ter uma 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 só estão disponíveis com CC ≥ 8.9 — na H100 e na RTX 4090; nas demais placas, a biblioteca passa automaticamente para BF16.

Como instalar o Transformer Engine?

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