Métodos de Treinamento de Baixa Precisão

O Accelerate fornece integrações para treinar com métodos de menor precisão usando hardware suportado especificado por meio dos pacotes TransformersEngine e torchao. Esta documentação ajudará a orientá-lo sobre qual hardware é suportado, como configurar seu Accelerator para aproveitar os métodos de baixa precisão e o que você pode esperar durante o treinamento.

O que significa treinar em FP8

Para explorar mais os detalhes técnicos do treinamento em FP8 com PyTorch e Accelerate, confira o concept_guide sobre por que isso pode ser difícil. Mas, essencialmente, em vez de treinar em BF16, alguns (ou todos) os aspectos do treinamento de um modelo podem ser executados usando 8 bits em vez de 16. O desafio é fazer isso sem degradar o desempenho final.

Isso é habilitado apenas em hardware NVIDIA específico, a saber:

• Tudo após as placas de vídeo de consumo da série 3000 (como a 4090)
• Arquiteturas de GPU baseadas em Hopper (como as H100 e H200)

Isso resultará em uma redução do uso de memória (pois reduzimos pela metade a memória necessária para algumas partes do treinamento) e um aumento no throughput deve ser observado também para modelos maiores que podem substituir certas camadas por camadas habilitadas para FP8.

Configurando o Accelerator

Atualmente, são suportados dois backends ativamente mantidos para FP8 (TransformersEngine e torchao), cada um com capacidades e configurações diferentes.

Para usar qualquer um deles, a mesma API principal é utilizada. Basta passar mixed_precision="fp8" para o Accelerator, durante o accelerate config quando solicitado sobre precisão mista, ou como parte do seu arquivo config.yaml na chave mixed_precision:

from accelerate import Accelerator
accelerator = Accelerator(mixed_precision="fp8")

Para especificar um backend (e personalizar outras partes da configuração de precisão mista FP8), você pode utilizar uma das dataclasses RecipeKwargs, como utils.AORecipeKwargs ou utils.TERecipeKwargs; você também pode esclarecer isso no seu arquivo de configuração yaml/durante o accelerate launch. Recomendamos o uso de TransformersEngine ou torchao para novos projetos:

from accelerate import Accelerator
from accelerate.utils import TERecipeKwargs, AORecipeKwargs
# Use TransformersEngine
kwargs = [TERecipeKwargs()]
# Or to use torchao
# kwargs = [AORecipeKwargs()]
accelerator = Accelerator(mixed_precision="fp8", kwarg_handlers=kwargs)

mixed_precision: fp8
fp8_config:
  amax_compute_algo: max
  amax_history_len: 1024
  backend: TE
  fp8_format: HYBRID
  interval: 1
  margin: 0
  override_linear_precision: (false, false, false)
  use_autocast_during_eval: false

Configurando o TransformersEngine

TransformersEngine oferece muitas opções para personalizar como e quais cálculos em FP8 são executados. Uma lista completa dos argumentos suportados e o que eles significam está disponível na documentação da NVIDIA, no entanto, eles são reafirmados como parte da docstring de FP8KwargsHandler para sua conveniência.

Accelerate tenta definir padrões sensatos, mas explorar e ajustar os vários parâmetros por conta própria pode potencialmente levar a um melhor desempenho.

Para usá-lo, especifique backend="te" e modifique quaisquer argumentos que desejar como parte do seu manipulador de kwarg:

from accelerate import Accelerator
from accelerate.utils import TERecipeKwargs
kwargs = [TERecipeKwargs(...)]
accelerator = Accelerator(mixed_precision="fp8", kwargs_handlers=kwargs)

Ou durante accelerate launch via --fp8_backend=te .... Use accelerate launch --fp8_backend=te -h para ver os argumentos relevantes.

Da mesma forma, isso pode ser definido no seu config.yaml:

mixed_precision: fp8
fp8_config:
    amax_compute_algo: max
    amax_history_len: 1024
    backend: TE
    fp8_format: HYBRID
    interval: 1
    margin: 0
    override_linear_precision: (false, false, false)
    use_autocast_during_eval: false

Configurando o torchao

torchao é um backend FP8 hackável, orientado por PyTorch, que visa ser mais acessível do que os dois motores anteriores. Uma das principais diferenças do ao em relação aos dois anteriores é que, para estabilidade numérica, descobriu-se que é geralmente melhor manter as primeiras e últimas camadas do modelo na precisão regular (seja FP32 ou BF16), e quantizar as demais camadas para FP8. Como resultado, uma configuração para ao se parece um pouco diferente:

Nota: esta API é experimental e está sujeita a alterações

from accelerate import Accelerator
from accelerate.utils import AORecipeKwargs, TorchDynamoPlugin, FullyShardedDataParallelPlugin
from torchao.float8 import Float8LinearConfig

fsdp2_plugin = FullyShardedDataParallelPlugin(
  fsdp_version=2,
  cpu_ram_efficient_loading=False, # CPU RAM efficient loading CANNOT work with fp8 torchao
  fsdp_auto_wrap_policy="TRANSFORMER_BASED_WRAP",
)
dynamo_plugin = TorchDynamoPlugin(
  backend="inductor",
  use_regional_compilation=True,
)
fp8_config = Float8LinearConfig(
  enable_fsdp_float8_all_gather=True, # Use FP8 all_gather in FSDP2
  pad_inner_dim=True,
)
kwargs = [AORecipeKwargs(
  config=fp8_config
)]
accelerator = Accelerator(
  mixed_precision="fp8",
  fsdp_plugin=fsdp2_plugin,
  dynamo_plugin=dynamo_plugin,
  kwarg_handlers=kwargs,
)

Ou durante accelerate launch via --fp8_backend=ao .... Use accelerate launch --fp8_backend=ao -h para ver os argumentos relevantes.

Da mesma forma, isso pode ser definido em config.yaml:

mixed_precision: fp8
fsdp_config:
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_cpu_ram_efficient_loading: false
  fsdp_version: 2
fp8_config:
  backend: AO
  pad_inner_dim: true
  enable_fsdp_float8_all_gather: true
dynamo_config:
  dynamo_backend: INDUCTOR
  dynamo_use_regional_compilation: true

Para saber mais sobre os parâmetros específicos a serem usados, consulte o repositório oficial do torchao.

Exemplo Zoo

Temos exemplos que demonstram o treinamento com FP8, tanto com o accelerate quanto com sua implementação subjacente, disponíveis no repositório do accelerate. Atualmente, suportamos scripts que demonstram:

• GPU única
• Paralelismo de Dados Distribuídos (Multi-GPU)
• Paralelismo de Dados Totalmente Fragmentado
• DeepSpeed ZeRO 1 a 3

Saiba mais aqui

Leituras adicionais

Para saber mais sobre treinamento em FP8, consulte os seguintes recursos:

• Nosso guia de conceitos detalhando mais sobre TransformersEngine e torchao
• A documentação do transformers-engine
• A documentação do torchao