PipeOffload: Verbetering van de schaalbaarheid van pijplijnparallelisme met geheugenoptimalisatie

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

Samenvatting

Pipeline-parallelisme (PP) wordt veel gebruikt voor het trainen van grote taalmodelen (LLMs), maar de schaalbaarheid ervan wordt vaak beperkt door het hoge geheugengebruik van activaties naarmate het aantal in-flight microbatches toeneemt met de mate van PP. In dit artikel richten we ons op het aanpakken van deze uitdaging door gebruik te maken van de onderbenutte strategie van geheugenoffloading in PP. Door empirisch onderzoek ontdekken we dat in de meeste standaardconfiguraties minstens de helft, en mogelijk alle, activaties kunnen worden geoffload met verwaarloosbare overhead. In gevallen waar volledige offloading niet mogelijk is, introduceren we een nieuwe selectieve offloadstrategie die het piekgeheugengebruik van activaties op een beter-dan-lineaire manier vermindert. Bovendien integreren we geheugenoffloading met andere technieken om gezamenlijk de totale doorvoer en geheugenbeperkingen in overweging te nemen. Onze experimenten bewijzen dat het geheugengebruik van activaties per apparaat effectief afneemt met het totale aantal fasen, waardoor PP een sterker alternatief wordt dan TP, met een versnelling tot wel 19\% en zelfs een lager geheugengebruik. De implementatie is open-source beschikbaar op https://github.com/sail-sg/zero-bubble-pipeline-parallelism{deze url}.

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}.