Parallélisme de tenseur

Le parallélisme de tenseur est un type de parallélisme des modèles dans lequel des poids, des gradients et des états d’optimiseur spécifiques sont répartis entre les appareils. Contrairement au parallélisme de pipeline, qui maintient les poids individuels intacts, mais partitionne l’ensemble de poids, les gradients ou l’optimiseur entre les dispositifs, le parallélisme de tenseur partitionne les poids individuels. Cela implique généralement un calcul distribué d'opérations, de modules ou de couches spécifiques du modèle.

Le parallélisme de tenseur est nécessaire dans les cas où un seul paramètre consomme la plus grande partie de la mémoire GPU (par exemple, de grandes tables d'incorporation avec une grande taille de vocabulaire ou une couche softmax volumineuse avec un grand nombre de classes). Dans ce cas, le traitement de ce tenseur ou de cette opération de grande taille comme une unité atomique est inefficace et nuit à l'équilibre de la charge mémoire.

SMP v2 s’intègre à Transformer Engine pour l’implémentation du parallélisme de tenseur et s’exécute au-dessus des API PyTorch FSDP. Vous pouvez activer le parallélisme de tenseur PyTorch FSDP et SMP simultanément, et déterminer le meilleur parallélisme des modèles pour de meilleures performances.

En pratique, le parallélisme de tenseur est particulièrement utile dans les scénarios suivants.

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

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

Pour la configuration de référence permettant d’appliquer le parallélisme de tenseur à ces modèles, consultez Conseils de configuration.

Configuration du parallélisme de tenseur

Pour tensor_parallel_degree, vous sélectionnez une valeur pour le degré de parallélisme de tenseur. La valeur doit diviser de manière égale le nombre de GPU de votre cluster. Par exemple, pour partitionner votre modèle lorsque vous utilisez une instance avec 8 GPU, choisissez 2, 4 ou 8. Nous vous recommandons de commencer par un petit nombre, puis de l’augmenter progressivement jusqu’à ce que le modèle tienne dans la mémoire du GPU.

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 les paramètres tensor_parallel_degree et random_seed, consultez Paramètres de configuration des caractéristiques principales de SMP v2.

Configuration SMP

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

Dans le script d’entraînement

Initialisez avec torch.sagemaker.init() pour activer SMP v2 et encapsulez votre modèle avec l’API torch.sagemaker.transform.

import torch.sagemaker as tsm
tsm.init()

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

Enregistrement et chargement des points de contrôle du transformeur Hugging Face

Une fois que la bibliothèque SMP a transformé un modèle, elle modifie le dictionnaire d’état (state_dict) du modèle. Cela signifie que le modèle devient incompatible avec les fonctionnalités de points de contrôle d’origine du transformeur Hugging Face. Pour gérer cela, la bibliothèque SMP fournit des API pour enregistrer les points de contrôle d’un modèle transformé dans une représentation du transformeur Hugging Face, et l’API torch.sagemaker.transform pour charger un point de contrôle du modèle de transformeur Hugging Face à des fins de peaufinage.

Pour plus d’informations sur la sauvegarde des points de contrôle lors de l’utilisation de la caractéristique de parallélisme de tenseur SMP v2, consultez Points de contrôle à l’aide de la SMP.

Pour plus d’informations sur le peaufinage d’un modèle en appliquant la caractéristique de parallélisme de tenseur SMP v2, consultez Peaufinage.