Déchargement de l'activation

Généralement, la passe avant calcule les activations au niveau de chaque couche et les conserve en GPU mémoire jusqu'à la fin de la passe arrière pour la couche correspondante. Le fait de décharger ces tenseurs en CPU mémoire après le transfert et de les récupérer au GPU moment où ils sont nécessaires peut permettre d'économiser beaucoup GPU de mémoire. 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 lors du CPU retour en arrière. Cela entraîne une dégradation majeure des performances lors de l'utilisation du déchargement d'activation.

SMPLa 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 repartir en arrière GPU sur ces activations. 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 fonctionnalité de déchargement SMP d'activation 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 :

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

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 les paramètres sm_activation_offloading et activation_loading_horizon, consultez SMPparamètres de configuration des fonctionnalités principales de la v2.

SMPconfiguration

{
    "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)