PipeOffload: Улучшение масштабируемости конвейерного параллелизма за счет оптимизации использования памяти

PipeOffload: Improving Scalability of Pipeline Parallelism with Memory Optimization

March 3, 2025

Авторы: Xinyi Wan, Penghui Qi, Guangxing Huang, Jialin Li, Min Lin

cs.AI

Аннотация

Параллелизм по конвейеру (Pipeline Parallelism, PP) широко используется для обучения больших языковых моделей (LLM), однако его масштабируемость часто ограничивается высоким потреблением памяти для активаций, поскольку количество одновременно обрабатываемых микропакетов увеличивается с ростом степени PP. В данной работе мы сосредоточились на решении этой проблемы, используя недостаточно изученную стратегию выгрузки памяти в PP. В результате эмпирического исследования мы обнаружили, что в большинстве стандартных конфигураций по крайней мере половина, а потенциально и все активации могут быть выгружены с незначительными накладными расходами. В случаях, когда полная выгрузка невозможна, мы предлагаем новую стратегию избирательной выгрузки, которая снижает пиковое потребление памяти для активаций лучше, чем линейно. Кроме того, мы интегрируем выгрузку памяти с другими методами, чтобы совместно учитывать общую пропускную способность и ограничения памяти. Наши эксперименты подтверждают, что потребление памяти для активаций на одно устройство эффективно снижается с увеличением общего количества этапов, что делает PP более предпочтительной альтернативой, чем TP, обеспечивая ускорение до 19\% при еще меньшем потреблении памяти. Реализация доступна в открытом доступе по адресу https://github.com/sail-sg/zero-bubble-pipeline-parallelism{этот 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}.