Stratégies pour économiser la mémoire lors de l’apprentissage dans les réseaux neuronaux profonds

Abstract

L’intelligence artificielle est un domaine qui a reçu beaucoup d’attention récemment. Son succès est dû aux progrès du Deep Learning, un sous-domaine qui réunit des méthodes d’apprentissage automatique basées sur les réseaux neuronaux. Ces réseaux neuronaux ont prouvé leur efficacité pour résoudre des problèmes très complexes dans différents domaines. Cependant, leur efficacité pour résoudre des problèmes dépend d’un certain nombre de facteurs : l’architecture du modèle, sa taille, comment et où l’entraînement a été effectué. La plupart des études indiquent que les modèles les plus gros permettent d’obtenir une meilleure précision, mais ils sont également plus coûteux à entraîner. Les principaux défis sont liés à la puissance de calcul et à la mémoire restreinte des machines : si le modèle est trop grand, son apprentissage peut prendre beaucoup de temps (des jours, voire des mois) ou, dans le pire des cas, il peut même ne pas tenir en mémoire. Pendant l’apprentissage, il est nécessaire de stocker à la fois les poids (paramètres du modèle), les activations(données calculées intermédiaires) et les états de l’optimiseur.Cette situation offre plusieurs opportunités pour traiter les problèmes de mémoire, en fonction de leurorigine. L’apprentissage peut être distribué sur plusieurs ressources de la plate-forme de calcul, et différentes techniques de parallélisation offrent différentes manières de distribuer la mémoire. En outre, les structures de données qui restent inactives pendant une longue période peuvent être temporairement déchargées vers un espace de stockage plus important, avec la possibilité de les récupérer ultérieurement (stratégies de déchargement). Enfin, les activations qui sont calculées à chaque itération peuvent être supprimées et recalculées plusieurs fois au cours de celle-ci (stratégies de rematérialisation). Les stratégies pour économiser la mémoire induisent généralement un surcoût en temps par rapport à l’exécution directe, par conséquent des problèmes d’optimisation doivent être considérés afin de choisir la meilleure approche pour chaque stratégie. Dans ce manuscrit, nous formulons et analysons des problèmes d’optimisation en relation avec diverses méthodes visant à réduire la consommation de mémoire pendant le processus d’apprentissage. En particulier, nous nous concentrons sur les stratégies de rematérialisation, de déchargement d’activation et de parallélisme de modèles pipelinés et pour chacune d’entre elles, nous concevons les solutions optimales sous un ensemble d’hypothèses. Enfin, nous proposons un outil entièrement fonctionnel appelé rotor qui combine l’activation offloading et la rematérialisation et qui peut être utilisé pour l’entraînement de grands modèles avec une surcharge minimale dans PyTorch, des modèles qui autrement ne tiendraient pas dans la mémoire.