PipeOffload : Amélioration de l'évolutivité du parallélisme de pipeline grâce à l'optimisation de la mémoire

PipeOffload: Improving Scalability of Pipeline Parallelism with Memory Optimization

March 3, 2025

Auteurs: Xinyi Wan, Penghui Qi, Guangxing Huang, Jialin Li, Min Lin

cs.AI

Résumé

Le parallélisme par pipeline (PP) est largement utilisé pour l'entraînement de grands modèles de langage (LLMs), mais sa scalabilité est souvent limitée par la consommation élevée de mémoire d'activation, qui augmente avec le nombre de micro-lots en cours d'exécution en fonction du degré de PP. Dans cet article, nous nous concentrons sur la résolution de ce défi en exploitant la stratégie sous-utilisée de déchargement de mémoire dans le PP. Grâce à une étude empirique, nous découvrons que dans la majorité des configurations standard, au moins la moitié, et potentiellement la totalité, des activations peuvent être déchargées avec un surcoût négligeable. Dans les cas où un déchargement complet n'est pas possible, nous introduisons une nouvelle stratégie de déchargement sélectif qui réduit la mémoire d'activation de pointe de manière meilleure que linéaire. De plus, nous intégrons le déchargement de mémoire avec d'autres techniques pour considérer conjointement le débit global et les limitations de mémoire. Nos expériences prouvent que la mémoire d'activation par dispositif diminue effectivement avec le nombre total d'étapes, faisant du PP une alternative plus robuste que le TP, offrant jusqu'à une accélération de 19 % avec une consommation de mémoire encore plus faible. L'implémentation est open-source à l'adresse suivante : https://github.com/sail-sg/zero-bubble-pipeline-parallelism{ce lien}.

English

Pipeline parallelism (PP) is widely used for training large language models (LLMs), yet its scalability is often constrained by high activation memory consumption as the number of in-flight microbatches grows with the degree of PP. In this paper, we focus on addressing this challenge by leveraging the under-explored memory offload strategy in PP. With empirical study, we discover that in the majority of standard configurations, at least half, and potentially all, of the activations can be offloaded with negligible overhead. In the cases where full overload is not possible, we introduce a novel selective offload strategy that decreases peak activation memory in a better-than-linear manner. Furthermore, we integrate memory offload with other techniques to jointly consider overall throughput and memory limitation. Our experiments proves that the per-device activation memory effectively reduces with the total number of stages, making PP a stronger alternative than TP, offering up to a 19\% acceleration with even lower memory consumption. The implementation is open-sourced at https://github.com/sail-sg/zero-bubble-pipeline-parallelism{this url}.