TransformerEngine
v1.13
QuickStart | Instalação | Guia do usuário | Exemplos | FP8 Convergência | Integrações | Notas de liberação
O Transformer Engine (TE) é uma biblioteca para acelerar modelos de transformadores nas GPUs NVIDIA, incluindo o uso de precisão de ponto flutuante de 8 bits (FP8) nas GPUs de tremonha, para fornecer melhor desempenho com menor utilização da memória em treinamento e inferência. A TE fornece uma coleção de blocos de construção altamente otimizados para arquiteturas de transformadores populares e uma API automática semelhante a precisão mista que pode ser usada perfeitamente com seu código específico da estrutura. O TE também inclui uma API agnóstica de C ++ da estrutura que pode ser integrada a outras bibliotecas de aprendizado profundo para ativar o suporte FP8 para transformadores.
À medida que o número de parâmetros nos modelos de transformadores continua a crescer, o treinamento e a inferência por arquiteturas como BERT, GPT e T5 se tornam muito memória e intensiva em computação. A maioria das estruturas de aprendizado profundo treinam com FP32 por padrão. Isso não é essencial, no entanto, para obter precisão total para muitos modelos de aprendizado profundo. Usando o treinamento de precisão mista, que combina precisão única (FP32) com formato de menor precisão (por exemplo, FP16) ao treinar um modelo, resulta em acelerações significativas com diferenças mínimas de precisão em comparação com o treinamento de FP32. Com a Hopper GPU Architecture, foi introduzida precisão FP8, o que oferece desempenho aprimorado em relação ao FP16, sem degradação na precisão. Embora todas as principais estruturas de aprendizado profundo suportem o FP16, o suporte ao FP8 não está disponível atualmente em estruturas em estruturas.
O TE aborda o problema do suporte do FP8, fornecendo APIs que se integram às bibliotecas populares do Modelo de Linguagem Large (LLM). Ele fornece uma API Python que consiste em módulos para construir facilmente uma camada de transformador, bem como uma biblioteca-agnóstico-estrutura em C ++, incluindo estruturas e kernels necessários para o suporte a FP8. Os módulos fornecidos pela TE para manter internamente os fatores de escala e outros valores necessários para o treinamento FP8, simplificando bastante o treinamento de precisão mista para os usuários.
import torch
import transformer_engine . pytorch as te
from transformer_engine . common import recipe
# Set dimensions.
in_features = 768
out_features = 3072
hidden_size = 2048
# Initialize model and inputs.
model = te . Linear ( in_features , out_features , bias = True )
inp = torch . randn ( hidden_size , in_features , device = "cuda" )
# Create an FP8 recipe. Note: All input args are optional.
fp8_recipe = recipe . DelayedScaling ( margin = 0 , fp8_format = recipe . Format . E4M3 )
# Enable autocasting for the forward pass
with te . fp8_autocast ( enabled = True , fp8_recipe = fp8_recipe ):
out = model ( inp )
loss = out . sum ()
loss . backward () import flax
import jax
import jax . numpy as jnp
import transformer_engine . jax as te
import transformer_engine . jax . flax as te_flax
from transformer_engine . common import recipe
BATCH = 32
SEQLEN = 128
HIDDEN = 1024
# Initialize RNG and inputs.
rng = jax . random . PRNGKey ( 0 )
init_rng , data_rng = jax . random . split ( rng )
inp = jax . random . normal ( data_rng , [ BATCH , SEQLEN , HIDDEN ], jnp . float32 )
# Create an FP8 recipe. Note: All input args are optional.
fp8_recipe = recipe . DelayedScaling ( margin = 0 , fp8_format = recipe . Format . HYBRID )
# Enable autocasting for the forward pass
with te . fp8_autocast ( enabled = True , fp8_recipe = fp8_recipe ):
model = te_flax . DenseGeneral ( features = HIDDEN )
def loss_fn ( params , other_vars , inp ):
out = model . apply ({ 'params' : params , ** other_vars }, inp )
return jnp . mean ( out )
# Initialize models.
variables = model . init ( init_rng , inp )
other_variables , params = flax . core . pop ( variables , 'params' )
# Construct the forward and backward function
fwd_bwd_fn = jax . value_and_grad ( loss_fn , argnums = ( 0 , 1 ))
for _ in range ( 10 ):
loss , ( param_grads , other_grads ) = fwd_bwd_fn ( params , other_variables , inp )A maneira mais rápida de começar com o Transformer Engine é usando imagens do Docker no catálogo NVIDIA GPU Cloud (NGC). Por exemplo, para usar o contêiner NGC Pytorch interativamente,
docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:23.10-py3Onde 23.10 é a versão do contêiner. Por exemplo, 23.10 para o lançamento de outubro de 2023.
Para instalar a versão estável mais recente do Transformer Engine,
pip install git+https://github.com/NVIDIA/TransformerEngine.git@stableIsso detectará automaticamente se alguma estrutura de aprendizado profundo suportada for instalada e criará o suporte ao mecanismo do transformador para eles. Para especificar explicitamente as estruturas, defina a variável de ambiente nvte_framework como uma lista separada por vírgula (por exemplo, nvte_framework = Jax, pytorch, paddle).
Como alternativa, o pacote pode ser instalado diretamente do Pypi do Transformer Engine, por exemplo,
pip install transformer_engine[pytorch]Para obter as ligações de Python necessárias para o mecanismo de transformador, as estruturas necessárias devem ser explicitamente especificadas como dependências extras em uma lista separada por vírgula (por exemplo, [Jax, Pytorch, Paddle]). O Transformer Engine envia rodas para a biblioteca principal, bem como para as extensões Paddlepddle. As distribuições de origem são enviadas para as extensões JAX e Pytorch.
Veja o guia de instalação.
Release do motor do transformador v0.11.0 Adiciona suporte para flashattion-2 em pytorch para melhorar o desempenho.
É um problema conhecido que a compilação Flashattion-2 é intensiva em recursos e requer uma grande quantidade de RAM (consulte Bug), o que pode levar a erros de memória fora da memória durante a instalação do mecanismo do transformador. Tente definir max_jobs = 1 no ambiente para contornar o problema.
Observe que o NGC Pytorch 23.08+ contêineres inclui flashattion-2.
Em um esforço para unificar a definição e o uso da máscara de atenção nas três estruturas no mecanismo de transformador, a máscara de preenchimento mudou da verdadeira inclusão de significado da posição correspondente na atenção à exclusão dessa posição em nossa implementação de Pytorch. Como v1.7, todos os tipos de máscara de atenção seguem a mesma definição em que verdadeiro significa mascarar a posição correspondente e os meios falsos, incluindo essa posição no cálculo da atenção.
Um exemplo dessa mudança é,
# for a batch of 3 sequences where `a`s, `b`s and `c`s are the useful tokens
# and `0`s are the padding tokens,
[a, a, a, 0, 0,
b, b, 0, 0, 0,
c, c, c, c, 0]
# the padding mask for this batch before v1.7 is,
[ True, True, True, False, False,
True, True, False, False, False,
True, True, True, True, False]
# and for v1.7 onwards it should be,
[False, False, False, True, True,
False, False, True, True, True,
False, False, False, False, True]O FP8 foi testado extensivamente em diferentes arquiteturas e configurações de modelos e não encontramos diferença significativa entre as curvas de perda de treinamento FP8 e BF16. O FP8 também foi validado para precisão nas tarefas LLM downstream (por exemplo, Lambada e Wikitext). Abaixo estão exemplos de modelos testados para convergência em diferentes estruturas.
| Modelo | Estrutura | Fonte |
|---|---|---|
| T5-770M | JAX/T5X | https://github.com/nvidia/jax-toolbox/tree/main/rosetta/rosetta/projects/t5x#convergence-and-performance |
| MPT-1.3b | Compositor de mosaico | https://www.mosaicml.com/blog/coreweave-nvidia-h100-part-1 |
| GPT-5B | JAX/PAXML | https://github.com/nvidia/jax-toolbox/tree/main/rosetta/rosetta/projects/pax#h100-sults |
| GPT-5B | Nemo Framework | Disponível a pedido |
| Llama2-7b | Alibaba Pai | https://mp.weixin.qq.com/s/nqt0ukxlbxyh5031zbdebq |
| T5-11B | JAX/T5X | Disponível a pedido |
| MPT-13B | Compositor de mosaico | https://www.databricks.com/blog/turboarged-trening-timizing-databricks-mosaic-ai-stack-fp8 |
| GPT-22B | Nemo Framework | Disponível a pedido |
| Llama2-70b | Alibaba Pai | https://mp.weixin.qq.com/s/nqt0ukxlbxyh5031zbdebq |
| GPT-175B | JAX/PAXML | https://github.com/nvidia/jax-toolbox/tree/main/rosetta/rosetta/projects/pax#h100-sults |
O mecanismo de transformador foi integrado às estruturas populares LLM, como:
Congratulamo -nos com contribuições para o Transformer Engine! Para contribuir com o Transformer Engine e fazer solicitações de tração, siga as diretrizes descritas no guia contribuinte.rst.
