Méthodes d’entraînement à basse précision

Accelerate fournit des intégrations pour l’entraînement avec des méthodes de basse précision en utilisant le matériel pris en charge spécifié, via les packages TransformersEngine et torchao. Cette documentation vous guide à travers le matériel pris en charge, la configuration de votre Accelerator pour exploiter les méthodes de basse précision, et ce que vous pouvez attendre lors de l’entraînement.

Ce que signifie l'entraînement en FP8

Pour explorer les détails techniques de l’entraînement en FP8 avec PyTorch et Accelerate, consultez le concept_guide qui explique pourquoi cela peut être difficile. En substance, plutôt que d’entraîner en BF16, certains (ou tous) les aspects de l’entraînement d’un modèle peuvent être effectués en 8 bits au lieu de 16. Le défi consiste à le faire sans dégrader les performances finales.

Cette fonctionnalité n'est activée que sur du matériel NVIDIA spécifique, à savoir :

• Tout ce qui se situe après les cartes graphiques grand public de la série 3000 (comme la 4090)
• Les architectures de GPU basées sur Hopper (telles que H100 et H200)

Le résultat sera une réduction de la mémoire utilisée (nous avons divisé par deux la mémoire nécessaire pour certaines parties de l’entraînement) ainsi qu’une augmentation du débit, qui should être également observée pour les modèles plus volumineux pouvant remplacer certaines couches par des couches compatibles FP8.

Configurer l'accélérateur

Deux backends pour FP8, activement maintenus, sont actuellement pris en charge (TransformersEngine et torchao), chacun avec des capacités et des configurations différentes.

Pour utiliser l’un ou l’autre, la même API de base est utilisée. Il suffit de passer mixed_precision="fp8" soit à Accelerator, soit lors de accelerate config lorsque l’on vous demande au sujet de la précision mixte, soit dans votre fichier config.yaml sous la clé mixed_precision :

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

Pour spécifier un backend (et personnaliser d'autres aspects de la configuration de la précision mixte FP8), vous pouvez utiliser l'une des dataclasses RecipeKwargs telles que utils.AORecipeKwargs ou utils.TERecipeKwargs ; vous pouvez également le préciser dans votre fichier de configuration yaml/lors de accelerate launch. Nous recommandons d'utiliser TransformersEngine ou torchao pour les nouveaux projets :

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

Configurer TransformersEngine

TransformersEngine offre de nombreuses options pour personnaliser la manière et le contenu des calculs en FP8. La liste complète des arguments pris en charge et leur signification est disponible dans la documentation de NVIDIA, mais elle est également reprise dans la docstring de FP8KwargsHandler pour votre confort.

Accelerate tente de définir des valeurs par défaut judicieuses, mais explorer et ajuster vous-même les différents paramètres peut potentiellement conduire à de meilleures performances.

Pour l'utiliser, spécifiez backend="te" et modifiez les arguments de votre choix dans votre gestionnaire de kwargs :

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

Ou lors de accelerate launch via --fp8_backend=te .... Utilisez accelerate launch --fp8_backend=te -h pour voir les arguments pertinents.

De même, cela peut être défini dans votre 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

Configurer torchao

torchao est un backend FP8 modifiable, piloté par PyTorch, conçu pour être plus accessible que les deux moteurs précédents. L’une des principales différences de ao par rapport aux deux précédents est que, pour la stabilité numérique, il est généralement préférable de conserver la première et la dernière couche du modèle à la précision standard (FP32 ou BF16), puis de quantiser les autres couches en FP8. En conséquence, une configuration pour ao se présente légèrement différemment :

Note : cette API est expérimentale et peut être modifiée

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 lors de accelerate launch via --fp8_backend=ao .... Utilisez accelerate launch --fp8_backend=ao -h pour voir les arguments pertinents.

De même, cela peut être défini dans 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

Pour en savoir plus sur les paramètres spécifiques à utiliser, veuillez consulter le dépôt officiel torchao.

Exemple Zoo

Nous disposons d'exemples illustrant l'entraînement avec FP8, à la fois avec accelerate et son implémentation sous-jacente, disponibles dans le dépôt accelerate. Nous prenons actuellement en charge des scripts illustrant :

• GPU unique
• Parallélisme de données distribué (multi-GPU)
• Parallélisme de données entièrement fragmentées
• DeepSpeed ZeRO 1 à 3

En savoir plus ici

Lecture complémentaire

Pour en savoir plus sur l'entraînement en FP8, consultez les ressources suivantes :

• Notre guide conceptuel qui détaille davantage TransformersEngine et torchao
• La documentation de transformers-engine
• La documentation de torchao