PipeOffload: Improving Scalability of Pipeline Parallelism with Memory
  Optimization
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}.
67c7b5900b05ab9c7e805433
Xinyi Wan
claimed_verified
2025-04-08T06:58:05.972Z
63510eea0b94548566dad923
/avatars/629eaaf810718259bf7588dc2e6cc0d5.svg
Xinyi Wan
user
ufotalent
67c7b5900b05ab9c7e805434
Penghui Qi
claimed_verified
2025-03-05T10:12:19.621Z
63885f1d0bebb233d8ad6e5b
https://cdn-avatars.huggingface.co/v1/production/uploads/1669881620925-noauth.jpeg
Penghui Qi
user
QPHutu
67c7b5900b05ab9c7e805435
Guangxing Huang
admin_assigned
2025-03-05T10:39:20.500Z
67c7b8c95d393f85737562e5
https://cdn-avatars.huggingface.co/v1/production/uploads/no-auth/KMmEVbXppNkfj1T9UOD7o.png
Guangxing Huang
user
huanggx-sea
67c7b5900b05ab9c7e805436
Jialin Li
admin_assigned
2025-03-05T10:39:29.789Z
658f147c50d39af7f4cd889d
/avatars/93bb33e975ab97c1f4179f20d384438f.svg
Jialin Li
user
JialinLi
67c7b5900b05ab9c7e805437
Min Lin
PipeOffload: Improving Scalability of Pipeline Parallelism with Memory
  Optimization
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}.
Memory offload in pipeline parallelism effectively reduces activation memory and improves throughput for large language models.
pipeline parallelism large language models activation memory memory offload selective offload strategy throughput memory limitation zero-bubble pipeline parallelism