Parallélisme de tenseur

Le parallélisme de tenseur est un type de parallélisme de modèle dans lequel des poids, des gradients et des états d'optimiseur spécifiques sont répartis entre les appareils. Contrairement au parallélisme des pipelines, qui permet de conserver les poids individuels intacts tout en répartissant l'ensemble des poids, des dégradés ou de l'optimiseur entre les appareils, le parallélisme tensoriel divise 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 la mise en œuvre du parallélisme des tenseurs et s'exécute au-dessus des API FSDP. PyTorch Vous pouvez activer simultanément le parallélisme des tenseurs PyTorch FSDP et SMP et déterminer le meilleur parallélisme du modèle pour de meilleures performances.

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

• Lorsque vous vous entraînez avec de longues durées de contexte, cela entraîne une mémoire d'activation élevée avec le FSDP uniquement.

• Lorsque vous vous entraînez avec de très grands clusters sur lesquels la taille globale du lot dépasse les limites souhaitées.

Modèles Hugging Face Transformer compatibles avec le parallélisme des tenseurs SMP

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

• GPT-Neox

• Lama 2

Pour la configuration de référence permettant d'appliquer le parallélisme des tenseurs à ces modèles, voir. Conseils de configuration

Configurer le parallélisme des tenseurs

Pourtensor_parallel_degree, vous sélectionnez une valeur pour le degré de parallélisme des tenseurs. La valeur doit diviser de manière égale le nombre de GPU de votre cluster. Par exemple, pour partager votre modèle lors de l'utilisation d'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 soit intégré à la mémoire du GPU.

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. Commencez avec la bibliothèque de parallélisme de 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 les paramètres tensor_parallel_degree et random_seed, consultez Paramètres de configuration des fonctionnalités principales du SMP v2.

Configuration du SMP

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

Dans votre script d'entraînement

Initialisez avec torch.sagemaker.init() pour activer SMP v2 et encapsulez votre modèle avec l'torch.sagemaker.transformAPI.

import torch.sagemaker as tsm
tsm.init()

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

Enregistrer et charger les points de contrôle du Hugging Face Transformer

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 point de contrôle d'origine du Hugging Face Transformer. Pour ce faire, la bibliothèque SMP fournit des API pour enregistrer les points de contrôle d'un modèle transformé dans une représentation de Hugging Face Transformer, et l'API pour charger un point de contrôle torch.sagemaker.transform du modèle Hugging Face Transformer à des fins de réglage précis.

Pour plus d'informations sur la sauvegarde des points de contrôle lors de l'utilisation de la fonction de parallélisme des tenseurs de SMP v2, consultez. Enregistrez et chargez des points de contrôle lors de l'utilisation du SMP

Pour plus d'informations sur le réglage précis d'un modèle en appliquant la fonction de parallélisme des tenseurs de SMP v2, consultez. Affinement