Parallélisme partitionné des données hybrides

Le parallélisme des données partitionnées est une technique d’entraînement distribué qui permet d’économiser de la mémoire et qui divise l’état d’un modèle (paramètres du modèle, gradients et états de l’optimiseur) entre les dispositifs. Cela vous permet de faire tenir un modèle plus grand ou d’augmenter la taille du lot en utilisant la mémoire GPU libérée. La bibliothèque SMP permet d'exécuter le parallélisme de données partitionné avec PyTorch Fully Sharded Data Parallel (FSDP). PyTorch Par défaut, le protocole FSDP partage l'ensemble des processeurs graphiques utilisés. Dans SMP v2, la bibliothèque propose ce parallélisme de données fragmenté en plus du PyTorch FSDP en étendant le sharding PyTorch hybride (HYBRID_SHARD), qui est l'une des stratégies de partitionnement proposées par FSDP :,,,. PyTorch FULL_SHARD SHARD_GRAD_OP HYBRID_SHARD _HYBRID_SHARD_ZERO2 L'extension du sharding hybride de cette manière permet de mettre en œuvre le sharding adapté à l'échelle, comme décrit dans le blog Scaling of Gigantic-model training on for Near-linear FSDP. AWS PyTorch

La bibliothèque SMP facilite son utilisation HYBRID_SHARD _HYBRID_SHARD_ZERO2 sur n'importe quel nombre configurable de GPU, en étendant le PyTorch FSDP natif qui prend en charge le partitionnement sur un seul nœud (HYBRID_SHARD) ou sur tous les GPU (). FULL_SHARD PyTorch Les appels FSDP peuvent rester tels quels, et il vous suffit d'ajouter l'hybrid_shard_degreeargument à la configuration SMP, comme indiqué dans l'exemple de code suivant. Il n'est pas nécessaire de modifier la valeur de l'sharding_strategyargument dans le wrapper PyTorch FSDP qui entoure votre modèle. PyTorch Vous pouvez passer ShardingStrategy.HYBRID_SHARD comme valeur. Autrement, la bibliothèque SMP remplace la stratégie du script et la définit sur ShardingStrategy.HYBRID_SHARD si vous spécifiez une valeur égale ou supérieure à 2 pour le paramètre hybrid_shard_degree.

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 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 hybrid_shard_degree, consultez Paramètres de configuration des caractéristiques principales de SMP v2.

Dictionnaire de configuration SMP

{ "hybrid_shard_degree": 16 }

Dans le script d’entraînement

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 needs to be created after FSDP wrapper
optimizer = ...