Déchargement de l'activation

Important

Dans SMP v2.2.0, la fonctionnalité d'activation et de déchargement de la bibliothèque SMP ne fonctionne pas. Utilisez plutôt le déchargement PyTorch d'activation natif.

Généralement, la passe directe calcule les activations au niveau de chaque couche et les conserve dans la mémoire du GPU jusqu'à la fin de la passe arrière pour la couche correspondante. Le fait de décharger ces tenseurs dans la mémoire du processeur après le transfert et de les récupérer sur le GPU lorsqu'ils sont nécessaires peut permettre d'économiser une utilisation substantielle de la mémoire du processeur graphique. PyTorch prend en charge les activations de déchargement, mais l'implémentation rend les GPU inactifs tandis que les activations sont récupérées du processeur lors du retour en arrière. Cela entraîne une dégradation majeure des performances lors de l'utilisation du déchargement d'activation.

SMP v2 améliore ce déchargement d'activation. Il prérécupère les activations à l'avance avant qu'elles ne soient nécessaires pour que le GPU commence à retransmettre ces activations. La fonction de pré-extraction permet d'exécuter les progrès de l'entraînement de manière plus efficace sans utiliser de GPU inactifs. Cela permet d'offrir les avantages d'une utilisation réduite de la mémoire sans dégradation des performances.

Vous pouvez conserver les PyTorch modules natifs pour décharger les activations dans votre script d'entraînement. Voici un exemple de structure d'application de la fonctionnalité de déchargement d'activation SMP dans votre script. Notez que le déchargement par activation n'est applicable que s'il est utilisé conjointement avecPoints de contrôle d'activation. Pour en savoir plus sur les outils de PyTorch point de contrôle natifs pour le déchargement des activations, voir :

• checkpoint_wrapper.py dans le PyTorch GitHubréférentiel

• Activation du point de contrôle PyTorch sur le blog Scaling Multimodal Foundation Models in TorchMultimodal with PyTorch Distributed.

Vous pouvez appliquer la fonction de déchargement d'activation SMP lors du point de contrôle d'PyTorch activation. Cela se fait en ajoutant les activation_loading_horizon paramètres sm_activation_offloading et au dictionnaire de configuration SMP pendantÉtape 2 : Lancer une offre de formation.

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 sm_activation_offloading et activation_loading_horizon, consultez Paramètres de configuration des fonctionnalités principales du SMP v2.

Configuration du SMP

{
    "activation_loading_horizon": 2,
    "sm_activation_offloading": True
}

Dans le script d'entraînement

Note

Lorsque vous activez la fonction de déchargement d'activation SMP, assurez-vous de l' PyTorch offload_wrapperutiliser également et de l'appliquer au module racine. La fonction de déchargement par activation SMP utilise le module racine pour déterminer à quel moment le transfert est effectué pour démarrer la préextraction.

import torch.sagemaker as tsm
tsm.init()

# Native PyTorch module for activation offloading
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
    apply_activation_checkpointing, 
    offload_wrapper,
)

model = FSDP(...)

# Activation offloading requires activation checkpointing.
apply_activation_checkpointing(
    model,
    check_fn=checkpoint_transformer_layers_policy,
)

model = offload_wrapper(model)