Parallélisme de contexte

Le parallélisme de contexte est un type de parallélisme des modèles 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 LayerNorm et RMSNorm, le parallélisme de contexte partitionne les entrées 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 de contexte et peut être utilisé conjointement avec le Parallélisme de tenseur PyTorch FSDP et SMP. Vous pouvez activer les trois parallélismes simultanément pour l’entraînement des modèles. Le parallélisme de contexte est bénéfique pour les modèles d’entraînement avec de grandes tailles d’activation et d’importantes longueurs de séquence. Il accélère le calcul des scores et des résultats de l’attention, en permettant à chaque appareil de calculer uniquement une partie des scores et des résultats le long de la dimension de séquence. Alors que le parallélisme de tenseur 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 de transformeurs Hugging Face compatibles avec le parallélisme de contexte SMP

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

Configuration du parallélisme de contexte

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

Les extraits de code suivants illustrent 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âche d’entraînement, tout en suivant le processus en deux étapes introduit dans Utilisation de la bibliothèque SageMaker de parallélisme des modèles v2. Vous n’avez pas besoin d’apporter de modifications à votre modèle PyTorch ni à la configuration PyTorch FSDP. Pour plus d’informations sur le paramètre context_parallel_degree, consultez Paramètres de configuration des caractéristiques principales de SMP v2.

Dans le 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’API torch.sagemaker.transform.

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

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

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