Métodos de entrenamiento de baja precisión

Accelerate ofrece integraciones para entrenar con métodos de menor precisión utilizando el hardware compatible especificado a través de los paquetes TransformersEngine y torchao. Esta documentación te guiará sobre qué hardware es compatible, cómo configurar tu Accelerator para aprovechar los métodos de baja precisión y qué puedes esperar durante el entrenamiento.

Qué significa el entrenamiento en FP8

Para explorar más detalles técnicos sobre el entrenamiento en FP8 con PyTorch y Accelerate, consulta la concept_guide sobre por qué esto puede ser difícil. En esencia, en lugar de entrenar en BF16, algunos (o todos) los aspectos del entrenamiento de un modelo pueden realizarse usando 8 bits en lugar de 16. El desafío es hacerlo sin degradar el rendimiento final.

Esto solo está habilitado en hardware NVIDIA específico, a saber:

• Todo lo que venga después de las tarjetas gráficas de consumo de la serie 3000 (como la 4090)
• Arquitecturas de GPU basadas en Hopper (como la H100 y la H200)

Esto resultará en una reducción del uso de memoria (ya que hemos reducido a la mitad la memoria necesaria para algunas partes del entrenamiento) y también se observará un aumento del rendimiento should para modelos más grandes que pueden reemplazar ciertas capas por capas habilitadas para FP8.

Configuración del acelerador

Actualmente se admiten dos backends para FP8 con mantenimiento activo (TransformersEngine y torchao), cada uno con capacidades y configuraciones diferentes.

Para usar cualquiera de ellas, se emplea la misma API principal. Simplemente pase mixed_precision="fp8" al Accelerator, durante accelerate config cuando se le pregunte sobre precisión mixta, o como parte de su archivo config.yaml en la clave mixed_precision:

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

Para especificar un backend (y personalizar otras partes de la configuración de precisión mixta FP8), puede utilizar una de las clases de datos RecipeKwargs, como utils.AORecipeKwargs o utils.TERecipeKwargs; también puede aclararlo en su configuración yaml/durante accelerate launch. Recomendamos usar TransformersEngine o torchao para nuevos proyectos:

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

Configuración de TransformersEngine

TransformersEngine ofrece muchas opciones para personalizar cómo y qué cálculos en FP8 se realizan. La lista completa de argumentos admitidos y su significado está disponible en la documentación de NVIDIA, aunque se repiten como parte de la docstring de FP8KwargsHandler para su comodidad.

Accelerate intenta establecer valores predeterminados razonables, pero explorar y ajustar los diversos parámetros por su cuenta puede conducir a un mejor rendimiento potencialmente.

Para usarlo, especifique backend="te" y modifique los argumentos que desee como parte de su manejador de kwarg:

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

O durante accelerate launch mediante --fp8_backend=te .... Use accelerate launch --fp8_backend=te -h para ver los argumentos relevantes.

De manera similar, esto se puede configurar en tu 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

Configuración de torchao

torchao es un backend FP8 modificable impulsado por PyTorch, que busca ser más accesible que los dos motores anteriores. Una de las diferencias clave de ao con respecto a los dos anteriores es que, por estabilidad numérica, se ha determinado que generalmente es mejor mantener las capas y la primera y la última del modelo en la precisión regular (ya sea FP32 o BF16), y luego cuantizar el resto de las capas a FP8. Como resultado, una configuración para ao se ve un poco diferente:

Nota: esta API es experimental y está sujeta a cambios

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,
)

O durante accelerate launch mediante --fp8_backend=ao .... Use accelerate launch --fp8_backend=ao -h para ver los argumentos relevantes.

De manera similar, esto se puede configurar en 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 obtener más información sobre los parámetros específicos que se deben utilizar, consulte el repositorio oficial de torchao.

Zoo de ejemplos

Tenemos ejemplos que muestran el entrenamiento con FP8 tanto con accelerate como con su implementación subyacente, disponibles en el repositorio de accelerate. Actualmente, los scripts que mostramos incluyen:

• GPU única
• Paralelismo de datos distribuido (multi-GPU)
• Paralelismo de datos totalmente fragmentado
• DeepSpeed ZeRO 1 a 3

Obtén más información aquí

Lecturas adicionales

Para obtener más información sobre el entrenamiento en FP8, consulte los siguientes recursos:

• Nuestra guía de conceptos con más detalles sobre TransformersEngine y torchao
• La documentación de transformers-engine
• La documentación de torchao