Memory Saving Strategies for Deep Neural Network Training
Stratégies pour économiser la mémoire lors de l’apprentissage dans les réseaux neuronaux profonds
Résumé
Artificial Intelligence is a field that has received a lot of attention recently. Its success is due to advances in Deep Learning, a sub-field that groups together machine learning methods based on neural networks. These neural networks have proven to be effective in solving very complex problems in different domains. However, their effectiveness depends on a number of factors: the architecture of the model, its size, how and where the training is performed... Most studies indicate that the large models are more likely to achieve the smallest error, but they are also more difficult to train. The main challenges are related to insufficient computational power and limited memory of the machines: if the model is too large then it can take a long time to be trained (days or even months), or it cannot even fit in memory in the worst case. During the training, it is necessary to store the weights (model parameters), the activations (intermediate computed data) and the optimizer states.This situation offers several opportunities to deal with memory problems, depending on their origin. Training can be distributed across multiple resources of the computing platform, and different parallelization techniques suggest different ways of dividing memory load. In addition, data structures that remain inactive for a long period of time can be temporarily offloaded to a larger storage space with the possibility of retrieving them later (offloading strategies). Furthermore, activations that are computed anew at each iteration can be deleted and recomputed several times within it (rematerialization strategies). Memory saving strategies usually induce a time overhead with respect to the direct execution, therefore optimization problems should be considered to choose the best approach for each strategy. In this manuscript, we formulate and analyze optimization problems in relation to various methods reducing memory consumption of the training process. In particular, we focus on rematerialization, activation offloading and pipelined model parallelism strategies, for each of them we design optimal solutions under a set of assumptions. Finally, we propose a fully functional tool called rotor that combines activation offloading and rematerialization and can be applied to training in PyTorch, allowing to process big models that otherwise would not fit into memory.
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.
Origine : Version validée par le jury (STAR)