Gradient checkpointing : comment économise-t-il la mémoire GPU ?
Traduction automatique
Cet article a été traduit automatiquement depuis la version originale en anglais.
Le gradient checkpointing conserve moins d’activations intermédiaires lors de la passe avant et recalcule celles qui manquent pendant le calcul arrière. Utilisez-le lorsque les activations empêchent le lot d’entraînement ou la longueur de séquence de tenir dans la mémoire GPU, et que le calcul supplémentaire est acceptable.
Il réduit la mémoire occupée par les activations conservées. Il ne supprime pas les poids du modèle, l’état de l’optimiseur ni tous les tampons temporaires.
Qu’est-ce qui est recalculé ?
Le calcul arrière a besoin des résultats intermédiaires de la passe avant pour calculer les gradients. Sans checkpointing, le framework conserve les résultats nécessaires. Avec le checkpointing, il enregistre certaines entrées ou certains résultats intermédiaires et réexécute les opérations concernées lorsque leurs valeurs sont nécessaires.
Par exemple, un framework peut enregistrer l’entrée d’un groupe de couches de transformer et supprimer leurs activations internes. Pendant le calcul arrière, il réexécute ce groupe à partir de l’entrée enregistrée. Cette entrée constitue un checkpoint d’activations, différent d’un checkpoint du modèle écrit sur disque pour permettre une reprise.
L’analyse originale de Chen et al. présente une stratégie à segments uniformes qui utilise O(sqrt(n)) de mémoire d’activations pour un réseau à n couches et nécessite environ une passe avant supplémentaire par étape d’entraînement. Les transformers réels ont des activations de tailles différentes, diverses implémentations de l’attention et des politiques propres aux frameworks. Mesurez la mémoire et le temps obtenus plutôt que de reprendre l’estimation asymptotique comme prévision de déploiement.
Activez l’intégration prévue par votre environnement
Pour les modèles pris en charge par Hugging Face Trainer, gradient_checkpointing=True active le mécanisme général. Avec FSDP, vérifiez plutôt l’intégration du checkpointing d’activations dans le framework. La documentation actuelle de Transformers Trainer recommande activation_checkpointing pour FSDP, car le gradient checkpointing générique peut entraîner un all-gather supplémentaire pendant le calcul arrière.
Comparez la même configuration d’entraînement avec et sans checkpointing. Relevez le pic de mémoire allouée sur GPU, le nombre d’exemples ou de tokens par seconde et vérifiez si la longueur de séquence prévue tient en mémoire. Vérifiez aussi que les gradients restent valides lorsque votre modèle utilise des opérations personnalisées ou l’entraînement d’adaptateurs.
Si l’état de l’optimiseur occupe la majeure partie de la mémoire, envisagez plutôt de répartir cet état entre les appareils ou d’utiliser un ajustement économe en paramètres ; recalculer les activations ne supprime pas cette allocation.
La section sur le gradient checkpointing du LLM Engineering Guide explique le calcul par segments.