PipeOffload: Miglioramento della Scalabilità del Parallelismo a Pipeline con Ottimizzazione della Memoria

PipeOffload: Improving Scalability of Pipeline Parallelism with Memory Optimization

March 3, 2025

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

cs.AI

Abstract

Il parallelismo pipeline (PP) è ampiamente utilizzato per l'addestramento di grandi modelli linguistici (LLM), ma la sua scalabilità è spesso limitata dall'elevato consumo di memoria delle attivazioni, poiché il numero di micro-batch in elaborazione cresce con il grado di PP. In questo articolo, ci concentriamo sull'affrontare questa sfida sfruttando la strategia di offload della memoria, ancora poco esplorata nel contesto del PP. Attraverso uno studio empirico, scopriamo che nella maggior parte delle configurazioni standard, almeno la metà, e potenzialmente tutte, le attivazioni possono essere offloadate con un overhead trascurabile. Nei casi in cui l'offload completo non è possibile, introduciamo una nuova strategia di offload selettivo che riduce la memoria di picco delle attivazioni in modo migliore rispetto a una riduzione lineare. Inoltre, integriamo l'offload della memoria con altre tecniche per considerare congiuntamente il throughput complessivo e i limiti di memoria. I nostri esperimenti dimostrano che la memoria delle attivazioni per dispositivo si riduce efficacemente con il numero totale di stadi, rendendo il PP un'alternativa più forte rispetto al TP, offrendo un'accelerazione fino al 19% con un consumo di memoria ancora inferiore. L'implementazione è open-source e disponibile all'indirizzo https://github.com/sail-sg/zero-bubble-pipeline-parallelism{questo 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}.