PipeOffload: Mejorando la Escalabilidad del Paralelismo en Tuberías con Optimización de Memoria

PipeOffload: Improving Scalability of Pipeline Parallelism with Memory Optimization

March 3, 2025

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

cs.AI

Resumen

El paralelismo de tuberías (PP) es ampliamente utilizado para entrenar modelos de lenguaje grandes (LLMs), sin embargo, su escalabilidad a menudo se ve limitada por el alto consumo de memoria de activaciones a medida que el número de micro-lotes en proceso aumenta con el grado de PP. En este artículo, nos enfocamos en abordar este desafío aprovechando la estrategia de descarga de memoria poco explorada en PP. Mediante un estudio empírico, descubrimos que en la mayoría de las configuraciones estándar, al menos la mitad, y potencialmente todas, las activaciones pueden ser descargadas con un impacto mínimo en el rendimiento. En los casos donde la descarga completa no es posible, introducimos una novedosa estrategia de descarga selectiva que reduce la memoria máxima de activaciones de manera mejor que lineal. Además, integramos la descarga de memoria con otras técnicas para considerar conjuntamente el rendimiento general y las limitaciones de memoria. Nuestros experimentos demuestran que la memoria de activaciones por dispositivo se reduce efectivamente con el número total de etapas, haciendo que PP sea una alternativa más fuerte que TP, ofreciendo una aceleración de hasta un 19\% con un consumo de memoria aún menor. La implementación es de código abierto en https://github.com/sail-sg/zero-bubble-pipeline-parallelism{este enlace}.

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