Treinamento de precisão mista

A biblioteca de paralelismo de SageMaker modelos (SMP) v2 oferece suporte a treinamento misto de precisão pronto para uso, integrando-se a estruturas de código aberto, como FSDP e Transformer Engine. PyTorch Para saber mais, consulte os tópicos a seguir.

Treinamento misto de precisão com FP8 instâncias P5 usando o Transformer Engine

A partir da biblioteca de paralelismo de SageMaker modelos (SMP) v2.2.0, a biblioteca SMP se integra ao Transformer Engine e oferece suporte a treinamento FP8 misto de precisão pronto para uso, mantendo a compatibilidade com o FSDP. PyTorch MixedPrecision Isso significa que você pode usar o PyTorch FSDP para treinamento de precisão mista e o Transformer Engine para treinamento. FP8 Para camadas de modelo não suportadas pelo recurso de FP8 treinamento do Transformer Engine, essas camadas retornam à precisão mista do PyTorch FSDP.

nota

O SMP v2 oferece FP8 suporte para os seguintes modelos de Hugging Face Transformer:

• GPT-NeoX (disponível no SMP v2.2.0 e versões posteriores)

• Llama 2 (disponível no SMP v2.2.0 e versões posteriores)

• Mixtral 8x7b e Mixtral 8x22b (disponíveis no SMP v2.5.0 e versões posteriores)

nota

Esse FP8 treinamento sobre o recurso P5 está disponível na seguinte combinação de bibliotecas de SageMaker e da PyTorch biblioteca:

• O SageMaker Python SDK v2.212.0 e versões posteriores

• PyTorch v2.2.0 e versões posteriores

FP8(precisão de ponto flutuante de 8 bits) é um tipo de dados que surgiu como outro paradigma para acelerar o treinamento de aprendizado profundo de modelos LLM. Com o lançamento da NVIDIA H100 GPUs suportando tipos de FP8 dados, você pode se beneficiar das vantagens das melhorias de desempenho nas instâncias P5 equipadas com o H100 GPUs, ao mesmo tempo em que acelera o treinamento distribuído com treinamento de precisão mista. FP8

O tipo FP8 de dados se ramifica ainda mais para os formatos E4M3 e E5M2. O E4M3 oferece uma melhor precisão, tem uma faixa dinâmica limitada e é ideal para o passe para frente no treinamento de modelos. O E5M2 tem uma faixa dinâmica mais ampla, mas com precisão reduzida, e é mais adequado para a passagem para trás, onde a precisão é menos crítica e uma faixa dinâmica mais ampla torna-se benéfica. Portanto, recomendamos que você use a receita da FP8 estratégia híbrida para aproveitar essas características de forma eficaz.

Para tipos de dados de meia precisão (FP16 e BF16), técnicas globais de escalonamento de perdas, como escalabilidade de perda estática ou escalabilidade dinâmica de perdas, lidam com problemas de convergência que surgem da perda de informações devido a gradientes de arredondamento em meia precisão. No entanto, a faixa dinâmica de FP8 é ainda mais estreita e as técnicas de escalonamento de perdas globais não são suficientes. Nesse ponto, precisamos de uma técnica de ajuste de escala por tensor mais refinada. O ajuste de escala atrasado é uma estratégia que seleciona um fator de escala com base nos valores absolutos máximos observados em vários tensores das iterações anteriores. Há uma desvantagem nessa estratégia: ela usa todos os benefícios de desempenho da FP8 computação, mas requer memória para manter o histórico de valores máximos dos tensores. Para saber mais sobre a estratégia de escalabilidade retardada em geral, consulte o artigo FP8 Formats for Deep Learning.

Na prática, o uso FP8 é útil em todos os cenários de treinamento em instâncias P5. É altamente recomendável ativar FP8 sempre que possível para melhorar o desempenho do treinamento.

O SMP v2 é compatível com o Transformer Engine pronto para uso. Portanto, ao executar o FP8 treinamento com SMP v2 em instâncias P5 de SageMaker AI (ml.p5.48xlarge), a única coisa que você precisa fazer é importar torch.sagemaker seu script de treinamento e continuar usando o pacote Python nativo do Transformer Engine. Para saber mais sobre como usar o Transformer Engine para FP8 treinamento em geral, consulte Usando FP8 com o Transformer Engine na documentação do NVIDIA Transformer Engine. O trecho de código a seguir mostra como devem ser as linhas de código para importar a biblioteca SMP e configurar seu script FP8 de treinamento.

import torch.sagemaker as tsm
import transformer_engine.pytorch as te
from transformer_engine.common.recipe import DelayedScaling, Format

# Initialize the SMP torch.sagemaker API.
tsm.init()

# Define a transformer model and wrap it with the torch.sagemaker.transform API.
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_config(ModelConfig)
model = tsm.transform(model)

# Enable E4M3 during forward pass, E5M2 during backward pass.
fp8_format = Format.HYBRID

# Create an FP8 recipe.
fp8_recipe = DelayedScaling(fp8_format=fp8_format, amax_history_len=32, amax_compute_algo="max")

# Enable FP8 autocasting.
with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe, fp8_group=tsm.state.world_process_group):
    out = model(inp)

loss = out.sum()
loss.backward()

Para encontrar um exemplo prático de FP8 treinamento com SMP v2 em instâncias P5, consulte o exemplo de caderno em Accelerate SageMaker PyTorch FSDP Training of LLama-v2 (ou GPT-Neox) com instâncias P5. FP8

Treinamento de precisão mista com tipos de dados de meia precisão usando PyTorch FSDP

O SMP v2 oferece suporte ao PyTorch FSDP MixedPrecision para trabalhos de treinamento em instâncias P4 e P5. PyTorch O FSDP fornece várias configurações para precisão mista para melhoria de desempenho e redução de memória.

nota

Esse treinamento misto de precisão com o recurso PyTorch FSDP está disponível na seguinte combinação de bibliotecas de SageMaker e da PyTorch biblioteca.

• SMP v2.0.0 e versões posteriores

• o SageMaker Python SDK v2.200.0 e versões posteriores

• PyTorch v2.0.1 e versões posteriores

A forma padrão de configurar um modelo para precisão mista é criar o modelo em float32 e, em seguida, permitir que o FSDP transmita os parâmetros para float16 ou bfloat16 dinamicamente ao passar uma política de MixedPrecision, conforme apresentado no trecho de código a seguir. Para obter mais informações sobre as opções de alteração de parâmetros, redução ou buffers para precisão mista PyTorch, consulte a MixedPrecisionAPI PyTorch FSDP na documentação. dtype PyTorch

# Native PyTorch API
from torch.distributed.fsdp import MixedPrecision

dtype = torch.bfloat16
mixed_precision_policy = MixedPrecision(
    param_dtype=dtype, reduce_dtype=dtype, buffer_dtype=dtype
)

model = FSDP(
    model,
    ...,
    mixed_precision=mixed_precision_policy
)

Observe que certos modelos (como o modelo Llama da Hugging Face Transformers) esperam buffers como float32. Para usar float32, substitua torch.bfloat16 por torch.float32, na linha que define o objeto dtype.