Parallélisme de données fragmenté hybride

Le parallélisme des données partitionnées est une technique d'entraînement distribuée économisant de la mémoire qui divise l'état d'un modèle (paramètres du modèle, dégradés et états de l'optimiseur) entre les appareils. Cela vous permet d'adapter un modèle plus grand ou d'augmenter la taille du lot en utilisant la mémoire libérée. GPU La SMP bibliothèque permet d'exécuter le parallélisme de données partitionné avec PyTorch Fully Sharded Data Parallel (). FSDP PyTorch FSDPpar défaut, des partitions sur l'ensemble des partitions GPUs utilisées. Dans la SMP version 2, la bibliothèque propose également ce parallélisme de données fragmentées PyTorch FSDP en étendant le sharding PyTorch hybride (HYBRID_SHARD), qui est l'une des stratégies de partitionnement fournies par PyTorch FSDP :,,,. FULL_SHARD SHARD_GRAD_OP HYBRID_SHARD _HYBRID_SHARD_ZERO2 L'extension du sharding hybride de cette manière permet de mettre en œuvre, scale-aware-sharding comme décrit dans le blog, une mise à l'échelle quasi linéaire d'un apprentissage de modèles gigantesques sur for. AWS PyTorch FSDP

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

Les extraits de code suivants montrent comment ajouter le module d'SMPinitialisation torch.sagemaker.init() à votre script de formation et configurer le dictionnaire de SMP configuration JSON au format 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 d'apporter des modifications à votre PyTorch modèle ou à votre PyTorch FSDPconfiguration. Pour plus d'informations sur le paramètre hybrid_shard_degree, consultez SMPparamètres de configuration des fonctionnalités principales de la v2.

SMPdictionnaire de configuration

{ "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 = ...