Déchargement d’activation

Généralement, la transmission vers l’avant calcule les activations au niveau de chaque couche et les conserve dans la mémoire du GPU jusqu’à la fin de la transmission vers l’arrière de 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 le déchargement des activations, mais l'implémentation les rend GPUs inactives pendant que les activations sont récupérées depuis le 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.

La bibliothèque SMP v2 améliore ce déchargement d’activation. Elle récupère les activations à l’avance avant qu’elles ne soient nécessaires pour que le GPU commence la transmission vers l’avant sur celles-ci. La fonction de prélecture permet d'exécuter les progrès de l'entraînement de manière plus efficace, sans interruption. GPUs 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 caractéristique de déchargement d’activation SMP dans votre script. Notez que le déchargement d’activation n’est applicable que s’il est utilisé conjointement avec Points 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 :

Vous pouvez appliquer la fonction de déchargement d'activation SMP lors du point de contrôle d'PyTorch activation. Pour ce faire, ajoutez les paramètres sm_activation_offloading et activation_loading_horizon au dictionnaire de configuration SMP pendant l’Étape 2 : lancer une tâche d’entraînement.

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. 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 caractéristiques principales de SMP v2.

Configuration SMP

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

Dans le script d’entraînement

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)