PipeOffload: Melhorando a Escalabilidade do Paralelismo em Pipeline com Otimização de Memória

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

Resumo

O paralelismo de pipeline (PP) é amplamente utilizado para o treinamento de grandes modelos de linguagem (LLMs), mas sua escalabilidade é frequentemente limitada pelo alto consumo de memória de ativação, à medida que o número de microbatches em execução aumenta com o grau de PP. Neste artigo, focamos em abordar esse desafio aproveitando a estratégia de descarregamento de memória pouco explorada no PP. Com estudos empíricos, descobrimos que, na maioria das configurações padrão, pelo menos metade, e potencialmente todas, as ativações podem ser descarregadas com sobrecarga insignificante. Nos casos em que o descarregamento completo não é possível, introduzimos uma nova estratégia de descarregamento seletivo que reduz o pico de memória de ativação de maneira melhor que linear. Além disso, integramos o descarregamento de memória com outras técnicas para considerar conjuntamente o rendimento geral e a limitação de memória. Nossos experimentos comprovam que a memória de ativação por dispositivo é efetivamente reduzida com o número total de estágios, tornando o PP uma alternativa mais forte que o TP, oferecendo uma aceleração de até 19% com consumo de memória ainda menor. A implementação é de código aberto em https://github.com/sail-sg/zero-bubble-pipeline-parallelism{este link}.

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