Entraînement de précision mixte

La bibliothèque de parallélisme des SageMaker modèles (SMP) v2 permet un entraînement de précision mixte prêt à l'emploi en s'intégrant à des frameworks open source tels que PyTorch FSDP et Transformer Engine. Pour en savoir plus, consultez les rubriques suivantes.

Entraînement de précision mixte avec FP8 sur des instances P5 à l'aide de Transformer Engine

À partir de la bibliothèque de parallélisme des SageMaker modèles (SMP) v2.2.0, la bibliothèque SMP s'intègre à Transformer Engine et prend en charge l'entraînement de précision mixte FP8 prêt à l'emploi, tout en préservant la compatibilité avec le FSDP. PyTorch MixedPrecision Cela signifie que vous pouvez utiliser à la fois le PyTorch FSDP pour l'entraînement de précision mixte et le Transformer Engine pour l'entraînement FP8. Pour les couches du modèle qui ne sont pas prises en charge par la fonction d'entraînement FP8 de Transformer Engine, ces couches reposent sur la précision PyTorch mixte FSDP.

Note

Le SMP v2 offre un support FP8 pour les modèles Hugging Face Transformer suivants :

• GPT-Neox

• Lama 2

Note

Cette formation FP8 sur la fonctionnalité P5 est disponible dans la combinaison suivante de bibliothèques de SageMaker et de PyTorch bibliothèque :

• SMP v2.2.0 et versions ultérieures

• le SDK SageMaker Python v2.212.0 et versions ultérieures

• PyTorch v2.2.0 et versions ultérieures

Le FP8 (précision à virgule flottante 8 bits) est un type de données qui est devenu un autre paradigme pour accélérer l'apprentissage en profondeur des modèles LLM. Avec le lancement des GPU NVIDIA H100 compatibles avec les types de données FP8, vous pouvez bénéficier des avantages liés à l'amélioration des performances des instances P5 équipées des GPU H100, tout en accélérant l'entraînement distribué grâce à l'entraînement de précision mixte FP8.

Le type de données FP8 se divise en outre vers les formats E4M3 et E5M2. L'E4M3 offre une meilleure précision, possède une plage dynamique limitée et est idéal pour la transmission vers l'avant lors de l'entraînement des modèles. L'E5M2 a une plage dynamique plus large, mais une précision réduite, et convient mieux à la passe arrière, où la précision est moins critique et où une plage dynamique plus large devient bénéfique. Par conséquent, nous vous recommandons d'utiliser la recette de stratégie hybride FP8 pour tirer parti de ces caractéristiques de manière efficace.

Pour les types de données de demi-précision (FP16 et BF16), les techniques globales de mise à l'échelle des pertes telles que la mise à l'échelle statique ou la mise à l'échelle dynamique des pertes permettent de résoudre les problèmes de convergence liés à la perte d'informations due à l'arrondissement des gradients en demi-précision. Cependant, la plage dynamique du FP8 est encore plus étroite et les techniques de mise à l'échelle globale des pertes ne sont pas suffisantes. À ce stade, nous avons besoin d'une technique de mise à l'échelle par tenseur plus fine. La mise à l'échelle différée est une stratégie qui sélectionne un facteur d'échelle basé sur les valeurs absolues maximales observées dans un certain nombre de tenseurs lors des itérations précédentes. Cette stratégie présente un inconvénient : elle utilise tous les avantages du calcul FP8 en termes de performances, mais nécessite de la mémoire pour conserver l'historique des valeurs maximales des tenseurs. Pour en savoir plus sur la stratégie de mise à l'échelle différée en général, consultez le document FP8 Formats for Deep Learning.

En pratique, l'utilisation de FP8 est utile dans tous les scénarios d'entraînement sur les instances P5. Nous recommandons vivement d'activer FP8 dans la mesure du possible pour améliorer les performances d'entraînement.

SMP v2 prend en charge Transformer Engine dès sa sortie de l'emballage. Par conséquent, lorsque vous exécutez un entraînement FP8 avec SMP v2 sur des instances P5 de SageMaker (ml.p5.48xlarge), la seule chose que vous devez faire est d'importer torch.sagemaker dans votre script d'entraînement et de continuer à utiliser le package Python natif de Transformer Engine. Pour en savoir plus sur l'utilisation de Transformer Engine pour l'entraînement FP8 en général, consultez la section Utilisation de FP8 avec Transformer Engine dans la documentation NVIDIA Transformer Engine. L'extrait de code suivant montre à quoi doivent ressembler les lignes de code permettant d'importer la bibliothèque SMP et de configurer FP8 dans votre script d'entraînement.

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

Pour trouver un exemple pratique d'entraînement FP8 avec SMP v2 sur des instances P5, consultez le bloc-notes d'exemples sur Accelerate SageMaker PyTorch FSDP Training of LLama-v2 (ou GPT-Neox) avec FP8 sur des instances P5.

Entraînement de précision mixte avec types de données de demi-précision à l'aide PyTorch du FSDP

SMP v2 prend en charge le PyTorch FSDP MixedPrecision pour les tâches de formation sur les instances P4 et P5. PyTorch Le FSDP propose différentes configurations pour une précision mixte, à la fois pour l'amélioration des performances et pour la réduction de la mémoire.

Note

Cet entraînement de précision mixte avec la fonction PyTorch FSDP est disponible dans la combinaison suivante de bibliothèques de SageMaker et de PyTorch bibliothèque.

• SMP v2.0.0 et versions ultérieures

• le SDK SageMaker Python v2.200.0 et versions ultérieures

• PyTorch v2.0.1 et versions ultérieures

La méthode standard pour configurer un modèle pour une précision mixte consiste à créer le modèle dansfloat32, puis à autoriser le FSDP à convertir les paramètres vers float16 ou à la bfloat16 volée en transmettant une MixedPrecision politique, comme indiqué dans l'extrait de code suivant. Pour plus d'informations sur les options permettant de modifier dtype les paramètres, la réduction ou les tampons pour une précision mixte PyTorch, consultez l'MixedPrecisionAPI PyTorch FSDP dans la documentation. 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
)

Notez que certains modèles (comme le modèle Hugging Face Transformers Llama) nécessitent des tampons tels que. float32 Pour l'utiliserfloat32, remplacez torch.bfloat16 par torch.float32 dans la ligne définissant l'dtypeobjet.