PipeOffload: Verbesserung der Skalierbarkeit von Pipeline-Parallelität durch Speicheroptimierung

PipeOffload: Improving Scalability of Pipeline Parallelism with Memory Optimization

March 3, 2025

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

cs.AI

Zusammenfassung

Pipeline-Parallelismus (PP) wird häufig für das Training großer Sprachmodelle (LLMs) eingesetzt, doch seine Skalierbarkeit wird oft durch den hohen Speicherverbrauch für Aktivierungen eingeschränkt, da die Anzahl der in Bearbeitung befindlichen Mikrobatches mit dem Grad des PP ansteigt. In diesem Beitrag konzentrieren wir uns darauf, diese Herausforderung durch die Nutzung der bisher wenig erforschten Speicherauslagerungsstrategie in PP zu bewältigen. Durch empirische Studien stellen wir fest, dass in den meisten Standardkonfigurationen mindestens die Hälfte, und potenziell alle, der Aktivierungen mit vernachlässigbarem Overhead ausgelagert werden können. In Fällen, in denen eine vollständige Auslagerung nicht möglich ist, führen wir eine neuartige selektive Auslagerungsstrategie ein, die den Spitzenverbrauch von Aktivierungsspeicher besser als linear reduziert. Darüber hinaus integrieren wir die Speicherauslagerung mit anderen Techniken, um den Gesamtdurchsatz und die Speicherbeschränkungen gemeinsam zu berücksichtigen. Unsere Experimente belegen, dass der Aktivierungsspeicher pro Gerät effektiv mit der Gesamtzahl der Stufen reduziert wird, wodurch PP eine stärkere Alternative zu TP darstellt und eine Beschleunigung von bis zu 19 % bei sogar geringerem Speicherverbrauch bietet. Die Implementierung ist unter https://github.com/sail-sg/zero-bubble-pipeline-parallelism{dieser URL} quelloffen verfügbar.

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