Parallélisme du contexte

Le parallélisme de contexte est un type de parallélisme de modèle qui répartit les activations du modèle selon la dimension de séquence. Contrairement aux autres techniques de parallélisme de séquence, qui ne font que partitionner le LayerNorm etRMSNorm, le parallélisme de contexte partitionne les entrées du réseau et toutes les activations intermédiaires le long de la dimension de séquence.

SMP v2 s'intègre à Transformer Engine pour le parallélisme du contexte et peut être utilisé conjointement avec PyTorch FSDP et SMP. Parallélisme de tenseur Vous pouvez activer les trois parallélismes simultanément pour l'entraînement des modèles. Le parallélisme du contexte est bénéfique pour les modèles d'entraînement dotés de grandes tailles d'activation et de longues longueurs de séquence. Il accélère le calcul des scores d'attention et des sorties d'attention, en permettant à chaque appareil de calculer uniquement une partie des scores et des sorties le long de la dimension de séquence. Alors que le parallélisme tensoriel accélère également le calcul en partitionnant le long de la dimension cachée, l'avantage du parallélisme de contexte est d'autant plus important que les exigences de calcul augmentent de façon quadratique avec la dimension de la séquence.

Modèles Hugging Face Transformer compatibles avec le parallélisme de contexte SMP

SMP v2 prend actuellement en charge le parallélisme contextuel pour les modèles de transformateurs Hugging Face suivants.

• GPT-Neox

• Llama 2 et Llama 3

• Mistral 7B

Configuration du parallélisme du contexte

Définissez une valeur entière pour le context_parallel_degree paramètre qui divise uniformément le nombre de GPUs dans votre cluster. Par exemple, si vous avez une instance à 8 GPU, utilisez 2, 4 ou 8 pourcontext_parallel_degree. Nous vous recommandons de commencer par une petite context_parallel_degree valeur et de l'augmenter progressivement jusqu'à ce que le modèle s'adapte à la mémoire du GPU avec la longueur de séquence d'entrée requise.

Les extraits de code suivants montrent comment ajouter le module d'initialisation SMP torch.sagemaker.init() à votre script d'entraînement et configurer le dictionnaire de configuration SMP au format JSON pour le lanceur de tâches de formation, tout en suivant le processus en deux étapes introduit dans. Utiliser la bibliothèque de parallélisme des SageMaker modèles v2 Il n'est pas nécessaire de modifier votre PyTorch modèle ou votre configuration PyTorch FSDP. Pour plus d’informations sur le paramètre context_parallel_degree, consultez Paramètres de configuration des fonctionnalités principales du SMP v2.

Dans votre script d'entraînement

Dans le cadre de l'étape 1, initialisez votre script avec torch.sagemaker.init() pour activer SMP v2 et encapsulez votre modèle avec l'torch.sagemaker.transformAPI.

À partir de SMP v2.6.0, vous pouvez utiliser l'argument cp_comm_type pour déterminer l'implémentation du parallélisme de contexte à utiliser. La bibliothèque SMP prend actuellement en charge deux implémentations : p2p et. all_gather L'p2pimplémentation utilise des appels d' peer-to-peerenvoi/réception pour l'accumulation de valeurs-clés lors de l'implémentation de l'attention et s'exécute de manière asynchrone, ce qui permet des chevauchements avec le calcul. all_gatherl'implémentation utilise plutôt l'opération AllGather collective et s'exécute de manière synchrone.

import torch.sagemaker as tsm
tsm.init()

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_config(..)
model = tsm.transform(model, cp_comm_type="p2p")

Configuration du SMP

Dans le cadre de l'étape 2, ajoutez le paramètre suivant au dictionnaire de configuration SMP pour l' SageMaker PyTorch estimateur.

{   
    ..., # other SMP config parameters
    "context_parallel_degree": 2
}