Le checkpointing de gradient est une technique utilisée dans le apprentissage profond pour réduire l'empreinte mémoire de l'entraînement des réseaux de neurones. Lors de la rétropropagation standard, un réseau doit stocker toutes les activations intermédiaires calculées lors de la passe avant pour calculer les gradients lors de la passe arrière. Pour des modèles très profonds, tels que les grands modèles de langage et les transformateurs, ce stockage peut dépasser la capacité mémoire du matériel disponible. Le checkpointing de gradient répond à ce problème en ne stockant pas chaque activation ; il conserve seulement un sous-ensemble et recalcule celles qui sont écartées à la demande lors de la passe arrière. Cela échange un coût de calcul accru contre une utilisation mémoire significativement réduite, permettant l'entraînement de modèles plus grands ou l'utilisation de tailles de lot plus importantes sur le même matériel.
La technique a été introduite en 2016 par des chercheurs de l'Carnegie Mellon University et d'OpenAI dans un article intitulé « Training Deep Nets with Sublinear Memory Cost ». Les auteurs, dont Tianqi Chen, Bing Xu, Chiyuan Zhang et Carlos Guestrin, ont démontré qu'en stockant les activations uniquement à certains points de contrôle (par exemple, toutes les quelques couches) et en recalculant le reste, le coût mémoire de l'entraînement d'un réseau profond pouvait être réduit de O(n) à O(sqrt(n)) pour un réseau à n couches, au prix d'environ une passe avant supplémentaire. Ce travail fondateur est depuis devenu un outil standard dans la communauté de l'apprentissage automatique, en particulier alors que la taille des modèles a considérablement augmenté.
Comment la rétropropagation standard utilise la mémoire
Dans une boucle d'entraînement conventionnelle, la passe avant calcule les activations pour chaque couche du réseau. Ces activations sont stockées en mémoire car la passe arrière en a besoin pour calculer les gradients via la règle de la chaîne. Pour un réseau avec L couches, cela nécessite de stocker L ensembles d'activations, chacun pouvant être volumineux. Par exemple, un réseau résiduel avec des centaines de couches ou un transformateur avec des dizaines de blocs d'attention peut accumuler des gigaoctets de données d'activation pour un seul exemple d'entraînement. Lors de l'entraînement avec de grandes tailles de lots, l'exigence mémoire évolue linéairement avec la taille du lot, devenant souvent le goulot d'étranglement principal.
La stratégie de checkpointing
Le checkpointing de gradient divise le réseau en segments, avec un point de contrôle à la frontière de chaque segment. Pendant la passe avant, seules les activations à ces frontières de points de contrôle sont sauvegardées en mémoire. Toutes les autres activations intermédiaires au sein d'un segment sont éliminées. Lorsque la passe arrière atteint un segment, elle recalcule la passe avant pour ce segment en utilisant l'activation de point de contrôle sauvegardée, régénérant ainsi les activations intermédiaires nécessaires au calcul du gradient. Cette recalcul ajoute une surcharge de calcul, généralement équivalente à une passe avant supplémentaire par étape d'entraînement, mais elle réduit considérablement l'utilisation de mémoire de pointe.
Le choix du placement des points de contrôle est un compromis. Plus de points de contrôle signifient moins de recalculs mais une utilisation mémoire plus élevée ; moins de points de contrôle signifient une mémoire plus faible mais plus de calcul. Le nombre optimal de points de contrôle pour un réseau avec n couches est approximativement sqrt(n), ce qui équilibre mémoire et calcul. En pratique, des frameworks comme PyTorch et TensorFlow permettent aux utilisateurs de spécifier des intervalles de points de contrôle ou d'utiliser des heuristiques automatiques.
Variantes et améliorations
Plusieurs raffinements de la technique originale ont été développés. Une variante courante est le checkpointing sélectif, où seuls certains types de couches (comme les blocs d'attention ou les couches convolutionnelles) sont sauvegardés, tandis que les autres sont stockés normalement. Une autre approche, appelée checkpointing de gradient économe en mémoire, utilise un schéma plus sophistiqué qui stocke les activations à plusieurs niveaux de granularité, réduisant encore la mémoire au prix de recalculs supplémentaires. Certaines implémentations proposent également un « déchargement », où les points de contrôle sont déplacés vers la mémoire du CPU ou le disque, bien que cela introduise une surcharge de transfert de données.
Dans le contexte des modèles Transformer (architecture), le checkpointing de gradient est souvent combiné avec d'autres techniques d'économie de mémoire comme l'écrêtage de gradient et l'entraînement en précision mixte. Par exemple, entraîner un modèle comme GPT-3, qui possède 175 milliards de paramètres, serait impossible sans de telles optimisations. La technique est également utilisée dans l'entraînement de grands modèles, où les économies de mémoire permettent aux praticiens de fonctionner sur un seul GPU au lieu d'un cluster.
Implémentation pratique
Dans les frameworks d'apprentissage profond modernes, le checkpointing de gradient est généralement exposé comme une fonctionnalité simple. Dans PyTorch, par exemple, le module torch.utils.checkpoint fournit une fonction checkpoint qui encapsule un module ou une séquence d'opérations. Lorsque le module encapsulé est exécuté, ses activations ne sont pas sauvegardées ; à la place, elles sont recalculées lors de la passe arrière. TensorFlow propose une fonctionnalité similaire via tf.recompute_grad. Ces implémentations gèrent la comptabilité automatiquement, permettant aux chercheurs d'adopter la technique sans modifier l'architecture de leur modèle.
Le surcoût de calcul du checkpointing de gradient n'est pas négligeable. Pour un réseau avec sqrt(n) points de contrôle, le calcul total de la passe avant pendant l'entraînement augmente d'environ 30 à 40 % par rapport à l'entraînement standard. Cependant, ce surcoût est souvent acceptable car l'alternative - réduire la taille du lot ou la taille du modèle - peut nuire à la convergence ou à la qualité du modèle. Dans de nombreux cas, le gain lié à l'utilisation d'une taille de lot plus importante compense le surcoût de recalcul.
Impact sur l'entraînement de grands modèles
Le checkpointing de gradient est devenu une pierre angulaire de l'entraînement de très grands modèles. Des entreprises comme OpenAI, Anthropic et Google DeepMind s'y appuient pour entraîner des modèles comportant des centaines de milliards de paramètres. Par exemple, entraîner un modèle de 70 milliards de paramètres sur un nœud unique avec 4 GPU nécessiterait de stocker des activations dépassant la mémoire combinée de ces GPU sans checkpointing. En utilisant le checkpointing de gradient, il est possible de faire tenir la charge d'entraînement dans le matériel disponible, bien qu'au prix de temps d'entraînement plus longs.
La technique est également essentielle pour l'IA générative avec de longues séquences, comme la génération de résumés de documents ou la traduction automatique. Dans ces cas, la mémoire d'activation augmente avec la longueur de la séquence, et le checkpointing permet d'utiliser des contextes plus longs sans dépasser les limites mémoire. Cela a directement permis le développement de modèles avec des fenêtres de contexte de 100 000 jetons ou plus.
Relation avec d'autres optimisations mémoire
Le checkpointing de gradient est souvent utilisé en conjonction avec d'autres techniques. Le batch normalization et le layer normalization ne réduisent pas directement la mémoire, mais améliorent la stabilité de l'entraînement, ce qui complète le checkpointing. La taille de modèle réduit le nombre de paramètres, mais les activations restent nécessaires, donc le checkpointing est toujours requis. L'augmentation de données augmente la taille effective du jeu de données mais n'affecte pas la mémoire d'activation. Dans l'entraînement distribué, le checkpointing de gradient peut être combiné avec la répartition des couches sur différents appareils pour réduire davantage la pression mémoire par appareil.
Une alternative notable est l'accumulation de gradient, qui simule une taille de lot plus grande en accumulant les gradients, mais elle réduit la mémoire pour les états de l'optimiseur et ne remplace pas le checkpointing. Une autre approche connexe est l'utilisation de couches réversibles, comme dans certains réseaux résiduels, où les activations peuvent être reconstruites à partir de la sortie, mais cela est moins général que le checkpointing.
Limites et compromis
La principale limite du checkpointing de gradient est le surcoût de calcul. Pour les modèles déjà limités par le calcul, l'entraînement peut ralentir de 20 à 40 %. De plus, cela ne réduit pas la mémoire nécessaire pour les paramètres du modèle ou les états de l'optimiseur, qui peuvent également être substantiels pour les grands modèles. Pour les modèles extrêmement grands, il peut être nécessaire de combiner le checkpointing avec d'autres techniques comme le déchargement des paramètres ou l'utilisation de matériel spécialisé avec une mémoire embarquée plus importante.
Un autre point subtil est que le recalcul peut introduire de légères différences numériques, bien que celles-ci soient généralement négligeables en pratique. Malgré ces défis, le checkpointing de gradient reste un outil largement utilisé et essentiel pour l'entraînement de modèles profonds.
Directions futures
À mesure que les modèles continuent de croître, de nouvelles stratégies de checkpointing émergent. Certaines recherches récentes utilisent des heuristiques apprises pour choisir les points de contrôle en fonction de la structure du modèle et du matériel. D'autres explorent le checkpointing au niveau des opérations individuelles ou combinent cette technique avec des approches comme la parallélisation par pipeline pour réduire davantage la pression mémoire. L'intérêt pour le checkpointing reste fort, car la mémoire demeure un facteur limitant dans le développement de modèles toujours plus grands et plus complexes.