La bibliothèque parallèle de modèles Amazon SageMaker accélère désormais les charges de travail PyTorch FSDP jusqu'à 20 % | Services Web Amazon

La formation sur les grands modèles de langage (LLM) a gagné en popularité au cours de l'année dernière avec la sortie de plusieurs modèles populaires tels que Llama 2, Falcon et Mistral. Les clients pré-entraînent et affinent désormais les LLM allant de 1 milliard à plus de 175 milliards de paramètres afin d'optimiser les performances des modèles pour les applications dans tous les secteurs, de la santé à la finance et au marketing.

Former des modèles performants à cette échelle peut être un défi. Les LLM très précis peuvent nécessiter des téraoctets de données de formation et des milliers, voire des millions d'heures de temps de calcul par accélérateur pour atteindre la précision cible. Pour terminer la formation et lancer les produits en temps opportun, les clients s'appuient sur des techniques de parallélisme pour répartir cette énorme charge de travail sur des milliers de dispositifs accélérateurs. Cependant, ces techniques de parallélisme peuvent être difficiles à utiliser : différentes techniques et bibliothèques ne sont compatibles qu'avec certaines charges de travail ou limitées à certaines architectures de modèles, les performances de formation peuvent être très sensibles à des configurations obscures et l'état de l'art évolue rapidement. En conséquence, les praticiens du machine learning doivent passer des semaines de préparation pour adapter leurs charges de travail LLM à de grands clusters de GPU.

Dans cet article, nous mettons en évidence les nouvelles fonctionnalités du Amazon Sage Maker bibliothèque de modèles parallèles (SMP) qui simplifie le processus de formation de grands modèles et vous aide à former les LLM plus rapidement. En particulier, nous couvrons la nouvelle expérience utilisateur simplifiée de la bibliothèque SMP qui s'appuie sur les API open source PyTorch Fully Sharded Data Parallel (FSDP), une fonctionnalité parallèle de tenseur étendue qui permet de former des modèles avec des centaines de milliards de paramètres et des optimisations de performances qui réduisent le temps de formation des modèles. et coûte jusqu'à 20 %.

Pour en savoir plus sur la bibliothèque parallèle de modèles SageMaker, reportez-vous à Bibliothèque de parallélisme de modèles SageMaker v2 Documentation. Vous pouvez également consulter notre exemples de cahiers pour commencer.

De nouvelles fonctionnalités qui simplifient et accélèrent la formation de grands modèles

Cet article traite des dernières fonctionnalités incluses dans la version v2.0 de la bibliothèque parallèle de modèles SageMaker. Ces fonctionnalités améliorent la convivialité de la bibliothèque, étendent les fonctionnalités et accélèrent la formation. Dans les sections suivantes, nous résumons les nouvelles fonctionnalités et expliquons comment vous pouvez utiliser la bibliothèque pour accélérer la formation de vos grands modèles.

Aligner SMP avec PyTorch open source

Depuis son lancement en 2020, SMP a permis une formation hautes performances à grande échelle sur les instances de calcul SageMaker. Avec cette dernière version majeure de SMP, la bibliothèque simplifie l'expérience utilisateur en alignant ses API avec PyTorch open source.

Offres PyTorch Parallélisme des données entièrement partagées (FSDP) comme méthode principale pour prendre en charge une charge de travail de formation importante sur de nombreux appareils informatiques. Comme le démontre l'extrait de code suivant, les API mises à jour de SMP pour des techniques telles que le parallélisme des données fragmentées reflètent celles de PyTorch. Vous pouvez simplement courir import torch.sagemaker et utilisez-le à la place de torch.

## training_script.py
import torch.sagemaker as tsm
tsm.init()

# Set up a PyTorch model
model = ...

# Wrap the PyTorch model using the PyTorch FSDP module
model = FSDP(
model,
...
)

optimizer = ...
...

Avec ces mises à jour des API de SMP, vous pouvez désormais bénéficier des avantages en termes de performances de SageMaker et de la bibliothèque SMP sans réviser vos scripts de formation PyTorch FSDP existants. Ce paradigme vous permet également d'utiliser la même base de code lors de la formation sur site que sur SageMaker, simplifiant ainsi l'expérience utilisateur pour les clients qui se forment dans plusieurs environnements.

Pour plus d'informations sur la façon d'activer SMP avec vos scripts de formation PyTorch FSDP existants, reportez-vous à Commencez avec SMP.

Intégration du parallélisme tensoriel pour permettre la formation sur des clusters massifs

Cette version de SMP étend également les capacités de PyTorch FSDP pour inclure des techniques de parallélisme tensoriel. L’un des problèmes liés à l’utilisation seule du parallélisme de données fragmentées est que vous pouvez rencontrer des problèmes de convergence lorsque vous augmentez la taille de votre cluster. En effet, les paramètres de partitionnement, les gradients et l'état de l'optimiseur dans les rangs parallèles de données augmentent également la taille globale de votre lot ; sur de grands clusters, cette taille de lot globale peut être poussée au-delà du seuil en dessous duquel le modèle convergerait. Vous devez incorporer une technique de parallélisme supplémentaire qui ne nécessite pas d'augmentation de la taille globale du lot à mesure que vous faites évoluer votre cluster.

Pour atténuer ce problème, SMP v2.0 introduit la possibilité de composer un parallélisme de données fragmentées avec un parallélisme tensoriel. Le parallélisme tensoriel permet d'augmenter la taille du cluster sans modifier la taille globale du lot ni affecter la convergence du modèle. Avec cette fonctionnalité, vous pouvez augmenter en toute sécurité le débit de formation en provisionnant des clusters avec 256 nœuds ou plus.

Aujourd'hui, le parallélisme tensoriel avec PyTorch FSDP n'est disponible qu'avec SMP v2. SMP v2 vous permet d'activer cette technique avec quelques lignes de changement de code et de débloquer une formation stable même sur de grands clusters. SMP v2 s'intègre à Moteur de transformateur pour son implémentation du parallélisme tensoriel et le rend compatible avec les API PyTorch FSDP. Vous pouvez activer simultanément le parallélisme des tenseurs PyTorch FSDP et SMP sans apporter de modifications à votre modèle PyTorch ou à la configuration PyTorch FSDP. Les extraits de code suivants montrent comment configurer le dictionnaire de configuration SMP au format JSON et ajouter le module d'initialisation SMP torch.sagemaker.init(), qui accepte le dictionnaire de configuration dans le backend lorsque vous démarrez la tâche de formation, à votre script de formation.

La configuration SMP est la suivante :

{
"tensor_parallel_degree": 8,
"tensor_parallel_seed": 0
}

Dans votre script de formation, utilisez le code suivant :

import torch.sagemaker as tsm
tsm.init()

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_config(..)
model = tsm.transform(model)

Pour en savoir plus sur l'utilisation du parallélisme tensoriel dans SMP, reportez-vous au parallélisme tenseur section de notre documentation.

Utilisez des fonctionnalités avancées pour accélérer la formation des modèles jusqu'à 20 %

En plus de permettre une formation distribuée sur des clusters comportant des centaines d'instances, SMP propose également des techniques d'optimisation qui peuvent accélérer la formation des modèles jusqu'à 20 %. Dans cette section, nous mettons en évidence quelques-unes de ces optimisations. Pour en savoir plus, reportez-vous au CARACTERISTIQUES de base section de notre documentation.

Partage hybride

Le parallélisme des données fragmentées est une technique de formation distribuée économisant la mémoire qui divise l'état d'un modèle (paramètres du modèle, gradients et états de l'optimiseur) entre les appareils. Cette empreinte mémoire plus petite vous permet d'intégrer un modèle plus grand dans votre cluster ou d'augmenter la taille du lot. Cependant, le parallélisme des données fragmentées augmente également les exigences de communication de votre tâche de formation, car les artefacts du modèle fragmenté sont fréquemment collectés à partir de différents appareils pendant la formation. De cette manière, le degré de partitionnement constitue une configuration importante qui compromet la consommation de mémoire et la surcharge de communication.

Par défaut, les fragments PyTorch FSDP modélisent les artefacts sur tous les périphériques accélérateurs de votre cluster. En fonction de votre travail de formation, cette méthode de partitionnement pourrait augmenter les frais de communication et créer un goulot d'étranglement. Pour vous aider, la bibliothèque SMP propose un parallélisme de données fragmentées hybrides configurable en plus de PyTorch FSDP. Cette fonctionnalité vous permet de définir le degré de partitionnement optimal pour votre charge de travail de formation. Spécifiez simplement le degré de partitionnement dans un objet JSON de configuration et incluez-le dans votre script de formation SMP.

La configuration SMP est la suivante :

{ "hybrid_shard_degree": 16 }

Pour en savoir plus sur les avantages du parallélisme des données fragmentées hybrides, reportez-vous à Mise à l'échelle quasi linéaire de la formation de modèles gigantesques sur AWS. Pour plus d'informations sur la mise en œuvre du partitionnement hybride avec votre script de formation FSDP existant, consultez parallélisme hybride des données partagées dans notre documentation.

Utiliser les opérations de communication collective SMDDP optimisées pour l'infrastructure AWS

Vous pouvez utiliser la bibliothèque SMP avec le Bibliothèque de parallélisme de données distribuées SageMaker (SMDDP) pour accélérer vos charges de travail de formation distribuées. SMDDP comprend un optimisé AllGather opération de communication collective conçue pour de meilleures performances sur les instances accélérées SageMaker p4d et p4de. Dans la formation distribuée, les opérations de communication collective sont utilisées pour synchroniser les informations entre les travailleurs GPU. AllGather est l'une des principales opérations de communication collective généralement utilisées dans le parallélisme des données fragmentées pour matérialiser les paramètres de couche avant les étapes de calcul avant et arrière. Pour les tâches de formation qui sont entravées par la communication, des opérations collectives plus rapides peuvent réduire le temps et les coûts de formation sans effets secondaires sur la convergence.

Pour utiliser la bibliothèque SMDDP, il vous suffit d'ajouter deux lignes de code à votre script de formation :

import torch.distributed as dist

# Initialize with SMDDP
import smdistributed.dataparallel.torch.torch_smddp
dist.init_process_group(backend="smddp") # Replacing "nccl"

# Initialize with SMP
import torch.sagemaker as tsm
tsm.init()

En plus de SMP, SMDDP prend en charge PyTorch FSDP et DeepSpeed ​​open source. Pour en savoir plus sur la bibliothèque SMDDP, consultez Exécutez une formation distribuée avec la bibliothèque de parallélisme de données distribuées SageMaker.

Déchargement d'activation

En règle générale, la passe avant de la formation du modèle calcule les activations à chaque couche et les conserve dans la mémoire GPU jusqu'à la fin de la passe arrière pour la couche correspondante. Ces activations stockées peuvent consommer une quantité importante de mémoire GPU pendant l'entraînement. Le déchargement d'activation est une technique qui déplace ces tenseurs vers la mémoire du processeur après la passe avant et les récupère ensuite vers le GPU lorsqu'ils sont nécessaires. Cette approche peut réduire considérablement l'utilisation de la mémoire GPU pendant l'entraînement.

Bien que PyTorch prenne en charge le déchargement des activations, sa mise en œuvre est inefficace et peut entraîner l'inactivité des GPU pendant que les activations sont récupérées depuis le CPU lors d'une passe arrière. Cela peut entraîner une dégradation significative des performances lors de l’utilisation du déchargement d’activation.

SMP v2 offre un algorithme de déchargement d'activation optimisé qui peut améliorer les performances d'entraînement. L'implémentation de SMP pré-extrait les activations avant qu'elles ne soient nécessaires sur le GPU, réduisant ainsi les temps d'inactivité.

Étant donné que SMP est construit sur les API de PyTorch, permettre un déchargement d'activation optimisé ne nécessite que quelques lignes de modification de code. Ajoutez simplement les configurations associées (sm_activation_offloading et de activation_loading_horizon paramètres) et incluez-les dans votre script de formation.

La configuration SMP est la suivante :

{
"activation_loading_horizon": 2,
"sm_activation_offloading": True
}

Dans le script de formation, utilisez le code suivant :

import torch.sagemaker as tsm
tsm.init()

# Native PyTorch module for activation offloading
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
apply_activation_checkpointing,
offload_wrapper,
)

model = FSDP(...)

# Activation offloading requires activation checkpointing.
apply_activation_checkpointing(
model,
check_fn=checkpoint_tformer_layers_policy,
)

model = offload_wrapper(model)

Pour en savoir plus sur les outils de point de contrôle open source PyTorch pour le déchargement d'activation, consultez le checkpoint_wrapper.py script dans le référentiel PyTorch GitHub et Points de contrôle d'activation dans le billet de blog PyTorch Mise à l'échelle des modèles de fondation multimodaux dans TorchMultimodal avec Pytorch Distributed. Pour en savoir plus sur la mise en œuvre optimisée du déchargement d'activation par SMP, consultez le déchargement d'activation section de notre documentation.

Au-delà du partitionnement hybride, du SMDDP et du déchargement d'activation, SMP offre des optimisations supplémentaires qui peuvent accélérer votre charge de travail de formation de modèles volumineux. Cela inclut des points de contrôle d’activation optimisés, une initialisation retardée des paramètres, etc. Pour en savoir plus, reportez-vous au CARACTERISTIQUES de base section de notre documentation.

Conclusion

À mesure que les ensembles de données, la taille des modèles et les clusters de formation continuent de croître, une formation distribuée efficace est de plus en plus essentielle pour une livraison rapide et abordable de modèles et de produits. La dernière version de la bibliothèque parallèle de modèles SageMaker vous aide à y parvenir en réduisant les modifications de code et en vous alignant sur les API PyTorch FSDP, permettant la formation sur des clusters massifs via le parallélisme tensoriel et des optimisations qui peuvent réduire le temps de formation jusqu'à 20 %.

Pour démarrer avec SMP v2, reportez-vous à notre Documentation et notre exemples de cahiers.

À propos des auteurs

Robert Van Dussen est chef de produit senior chez Amazon SageMaker. Il dirige les frameworks, les compilateurs et les techniques d'optimisation pour la formation en deep learning.

Luis Quintela est le responsable du développement logiciel pour la bibliothèque parallèle de modèles AWS SageMaker. Pendant son temps libre, on peut le trouver sur sa Harley dans la région de la baie de SF.

Gautam Kumar est ingénieur logiciel chez AWS AI Deep Learning. Il est passionné par la création d'outils et de systèmes pour l'IA. Dans ses temps libres, il aime faire du vélo et lire des livres.

Rahul Huilgol est ingénieur principal en développement logiciel dans le domaine de l'apprentissage profond distribué chez Amazon Web Services.

• Contenu propulsé par le référencement et distribution de relations publiques. Soyez amplifié aujourd'hui.
• PlatoData.Network Ai générative verticale. Autonomisez-vous. Accéder ici.
• PlatoAiStream. Intelligence Web3. Connaissance Amplifiée. Accéder ici.
• PlatonESG. Carbone, Technologie propre, Énergie, Environnement, Solaire, La gestion des déchets. Accéder ici.
• PlatoHealth. Veille biotechnologique et essais cliniques. Accéder ici.
• La source: https://aws.amazon.com/blogs/machine-learning/amazon-sagemaker-model-parallel-library-now-accelerates-pytorch-fsdp-workloads-by-up-to-20/