PipeOffload: Improving Scalability of Pipeline Parallelism with Memory Optimization

We thank the reviewer for the insightful feedbacks. See response to individual points below.

Whether the proposed strategy is suitable for other pipeline parallel strategy like DualPipe? The authors can present some analysis theoretically and empirically.

TL;DR: Yes, but different devices will have different memory deduction.

The selective offload strategy is applicable to bi-directional schedules like Chimera and DualPipe, as they involve multiple stages with different directions on the same devices. In these discussions, we refer to the stage on device $i$ with a forward pass from device $i$ to device $i+1$ as stage 0, and the reverse stage as stage 1.

In DualPipe, stages on the same device have varying lifespans. For instance, on device 0, stage 0 has a longer lifespan, while stage 1 has a negligible one. This creates an opportunity for better-than-linear memory reduction, especially since DualPipe uses a "Uniform Repeat" strategy, as mentioned in Section 2.

However, the lifespan distribution varies across devices. For example, stage $\frac{d}{2}$ has nearly equal lifespans for both stages, and selective offloading leads to linear memory reduction. In a rough analysis, offloading one stage on devices 0 and $d-1$ results in approximately 100% activation reduction, while on device $\frac{d}{2}$, it results in about 50%. The memory reduction for other devices can be linearly interpolated.

Finally, we emphasize that DualPipe's main advantage is in expert-parallel (EP) pipeline scheduling, where it fuses forward and backward passes to overlap EP communication with computation.

Whether the authors can provide the core codebase of the experiments? It may make the experiments reproducible.

We are going to open-source the code base for our implementation and experiments for a broader impact in the community. We did not upload it as supplementary material because we were not confident in anonymizing the code sufficiently to comply with double-blind review rules.

We thank the reviewer for the insightful feedbacks. See response to individual points below

The selective offload strategy introduces additional complexity in managing activation lifespans, which could complicate implementation and maintenance.

Every new feature requires some implementation, however, we try our best to make implementation simple. The offloading majorly leverages the existing PyTorch Tensor.to() function to do D2H and H2D transfers. The scheduling code handles offloading as PP native passes and uses simple heuristics to determine the ordering between them.

The effectiveness of the proposed strategies may vary across different hardware configurations, potentially limiting their general applicability.

We agree that the offload time ratio $k$ is dependent on hardware, particularly the host-to-device bandwidth. Our implementation allows for adjustable scheduling to accommodate different hardware configurations.

Additionally, mainstream GPUs are rapidly improving in PCIe bandwidth (e.g., A100: 64 GB/s, H100: 128 GB/s, B100: 256 GB/s, Grace Hopper: 900 GB/s). Given this trend, we expect our method to remain generalizable to future devices.

How seamlessly can PipeOffload be integrated into existing machine learning frameworks that utilize 3D parallelism

Our implementation is based on Megatron-LM, one of the most popular LLM training frameworks. PipeOffload is orthogonal to other parallelism schemes and requires minimal additional effort when integrated into a hybrid parallelism strategy.

is it feasible to apply PipeOffload in conjunction with 1F1B-I to mitigate this issue?

Yes, the same method can be applied to 1F1B-I. However, as discussed in Section 2 of our paper, selective offloading in 1F1B-I only achieves linear memory savings. This setup is still preferable when higher throughput is the priority.

in Figure 7 of your paper, how did you determine the durations for offloading/reloading and forward/backward computations? Could you elaborate on the methodology used to decide these lengths?

In our current implementation, for simplicity, we use the ratio $k$ between offload time and compute time, as calculated by Formula 1, while assuming F/B/W passes run for identical time. Experiments show that this approach works well. However, we anticipate that using more accurately profiled timings for compute and offload passes will lead to even better performance.

Can you elaborate on the memory imbalance across different stages? Does PipeOffload need to offload/reload the same amount of activations for different stages?

Could you clarify whether this imbalance refers to the different activation memory for different stages on the same device, or different stages on different devices? If it's about different stages on the same device, the imbalance arises from the varying lifespans of activations across stages. For example, in the forward pass, stage 0 computes before stage 1, while the backward pass is computed after stage 1. This results in stage 0’s activations living longer than those of stage 1, contributing to higher peak memory for stage 0. If multiple stages are offloaded, the number of activations to be offloaded will be the same across stages, but the peak memory reduction will differ, with stage 0 providing a greater reduction than stage 1.

How does PipeOffload's performance compare to that of bidirectional pipeline parallelism? Is it possible to combine PipeOffload with bidirectional pipeline parallelism strategies?

The same insight of selective offloading can also be applied to bidirectional PP, as in bidirectional PP there're also multiple stages on the same device.

How to decide the offload schedule for different topologies?

The topology influences how offload/reload passes are co-scheduled across different devices. For instance, if two GPUs are connected to the same PCIe switch, we schedule them in a "Synchronized Interleaved" pattern, as described in Table 3. If the GPUs have independent PCIe connections, we schedule offload/reload passes independently on each device.

We thank the reviewer for the insightful feedbacks. Please see individual points below.

Including learning curves could enhance the understanding of convergence and correctness.

We checked that our methods produces identical loss curves compared to vanilla 1F1B-I. We'll include this in the PDF. Convergence Graph

Incorporating comparisons with Chimera and DualPipe

Our primary contribution is memory savings with minimal overhead. Since DualPipe, Chimera, and 1F1B share identical peak memory consumption, comparing against 1F1B and 1F1B-I sufficiently supports our claim.

Additionally, DualPipe's advantage lies in expert-parallel (EP) pipeline scheduling, where it fuses forward and backward passes to overlap EP communication with computation. As our work does not focus on EP integration, a comparison with DualPipe with EP is beyond our scope.

Moreover, we have added experiments comparing our approach with state-of-the-art memory-saving methods, including the V-Shape PP schedules from [2] and 1F1B-I with rematerialization (recompute). The results, shown in the image below, demonstrate that our methods push the Pareto frontier of the throughput-memory trade-off.

Pareto Frontier Figure

In this figure, 1F1B-I-R represents Interleaved 1F1B with full rematerialization, while V-Min and V-Half are from [2].

It is unclear how bandwidth affects the performance of the proposed method. What is the minimum memory bandwidth required to hide the offload overhead? Is this method applicable on systems with only PCIe?

Common hardware, including PCIe-only systems, already provides sufficient bandwidth for activation offloading. Our method remains effective under these conditions.

We extensively analyzed this in the introduction section. Formula 1 and Figure 1 demonstrate that typical hardware can offload a large portion (at least half, possibly all) of activations. In Formula 1, $k$ describes the ratio between offload time and compute time, and we see an important term $B_o$ which is the offload bandwidth. We plotted practical $k$ values based on A100 GPUs with PCIe v4 on Figure 1 (left) and have shown that a $k$ is actually small. Section 4.3 further discusses $k$ across different hardware, confirming similar trends.

Notably this only relies on the host-device PCIe bandwidth, so it also works on systems with only PCIe.

To clarify this emphasis on bandwidth, we will update the PDF to highlight the role of Formula 1 and Figure 1.

it is unclear in what case the proposed pipeline parallelism method can outperform tensor parallelism. Discussion on being integrated with other parallelism strategies

Thanks for raising this important discussion on TP vs. PP. In short, PP should be preferred when memory allows, as it generally provides higher throughput.

Our results in Figure 11 align with prior work [1][2], showing PP outperforms TP due to TP’s high communication overhead. However, when memory is limited, TP can still help reduce activation memory. This highlights our contribution—reducing PP memory usage so that less TP is needed.

On implementation, PipeOffload is orthogonal to other parallelization strategies and can be integrated into hybrid parallelism.

A summary of a detailed analysis of TP vs PP that we'll add to the PDF:

• Activation Memory: When using $d$ devices, For TP, each GPU has approximately $\frac{1}{d}$ of the activation memory of the entire model. For our PipeOffload, with selective offloading each device consumes approximately $\frac{1}{4}$ for a selective offload and $O(\frac{1}{vd})$ for a fully offload. Considering in practical scenarios $d\le8$ for TP, PipeOffload has a similar, or even situationally (when $k\lt1$, see Figure 1 in our paper) lower per-device activation memory compared to TP.
• Throughput: Prior studies [2][3] show that TP with 8 devices incurs a 20–30% throughput overhead. For PipeOffload, the bubble rate is $\lt\frac{d-1}{n+d-1}$, where n is the number of microbatches. In most practical cases, the bubble rate of PipeOffload is lower than the overhead of TP.

Other methods like 1F1B-I may achieve better MFU if sufficient memory is available.

We agree that when memory is sufficient, throughput should be prioritized. However, in many practical cases, "sufficient memory" is only achieved by using TP, which incurs significant throughput overhead, or by applying activation rematerialization. Our work provides an effective alternative for acceleration in such scenarios.

References

[1] BPipe: Memory-Balanced Pipeline Parallelism for Training Large Language Models

[2] Pipeline Parallelism with Controllable Memory

[3] Megatron-LM: Training Multi-Billion Parameter Language Models using GPU Model Parallelism.

We thank the reviewer for the valuable comments. We respond to individual points from your review below.

The paper lacks clarity regarding its novel contributions ... the distinction between established and novel concepts, such as Selective Offload and Zero-Bubble Strategy, is not adequately delineated ... The design section could benefit from clarifying the originality ...

Our main contributions are:

• For the first time, we identify the oppotunity that all activation memory of PP can be offloaded with negligible overhead, which is critical but overlooked for the scalability of PP.
• We propose a novel selective offload strategy in PP, which can achieve better-than-linear memory reduction. To the best of our knowledge, this is the first selective strategy for offload.
• We extend the popular interleaved 1F1B into a generalized form for the first time, which can smoothly trade memory with minimal throughput loss, providing a flexible solution that can be tailored to specific system requirements.
• We provide implementation details in Section 4 and Appendix, which is critical to achieve high performance for offload in PP.

As referenced in the paper, Zero-Bubble was proposed by [1]. We apply this strategy as a preliminary step in our generalized interleaving strategy (Section 3.2). We don't claim it as our contribution, as it's trivial.

It would be helpful to address the fundamental differences against the existing work [2]

Although we are inspired by some concepts (building-block, lifespan) from [2], our work has fundamentally different contributions and insights compared to [2]. Specifically, we focus on selective offloading (Section 2), designing a novel scheduling methodology to facilitate offloading (Section 3), and implementing optimized offload techniques (Section 4). These contributions are centered on offloading, whereas [2] aims to reduce activation memory through schedule-based methods, and does not mention offloading at all.

Is GIS(-H) an existing method or an entirely new design? ...

The generalized interleaving strategy (Section 3.2) is entirely new. Both GIS and GIS-H can be seen as special cases of tuning $g$ to $g=d$ or $g=\frac{d}{2}$. While GIS happens to resemble Figure 17$c$ in [2], it is not identical. To the best of our knowledge, GIS-H ($g=\frac{d}{2}$) and all GIS schedules with $g \lt d$ are entirely novel.

acceleration factor is only measured against TP instead of other PP methods. what is the acceleration of PipeOffload compared to SOTA PP methods (rather than TP ones)? Given PP being a research hotspot, expanding the comparison baselines to include methods like BPipe, GPipe, and PipeFill ... if any of them is open-source.

Our primary focus is memory reduction to expand the applicability of pipeline parallelism (PP). Like all memory-saving methods, ours does not improve PP throughput; rather, it trades throughput for memory savings. This makes comparisons with tensor parallelism (TP) relevant, as TP is often preferred due to PP’s high memory requirements.

BPipe and [2] are important works sharing the same target of memory saving with us, while only [2] is open-source. We conducted experiments comparing our approach with [2] and 1F1B-I with rematerialization which are the SoTA memory saving PP schedule. The results, shown in the image below, demonstrate that our methods (PipeOffload and GIS schedules) push the Pareto frontier of memory and throughput.

Pareto Frontier Figure

In this figure, 1F1B-I-R represents Interleaved 1F1B with full rematerialization, while V-Min and V-Half are from [2].

Regarding Figure 2, what is the cause of the asymmetry between Stages 0 and 1?

Are you asking why the trends of Stage 0 and Stage 1 are reversed? Intuitively, after merging both stages, the final schedule follows a one-forward-one-backward pattern, ensuring a constant total activation memory across the two stages. As a result, when the memory usage of one stage increases, the other decreases.

Regarding Figure 3, how is the curve modeled and calculated? What does the kink near x=13 mean?

We measure the percentage of memory reduction at $x \in \{0, 1, \ldots, 16\}$ and connect these points to form a curve. The kink near $x=13$ is simply because the evaluation points are discrete ($x$ must be an integer).

... It would be advantageous to also cross-validate certain equations, for example, memory and bubble amount ...

As mentioned in the caption of Figure 8, we provide detailed numbers in Figure 13.

It would be beneficial to explain why Megatron-LM is selected as the evaluation target if that is not common knowledge.

We believe it's a common knowledge that Megatron-LM is SOTA and a common choice.

[1] Qi, P., Wan, X., Huang, G., and Lin, M. Zero bubble pipeline parallelism.

[2] Qi, P., Wan, X., Amar, N., and Lin, M. Pipeline parallelism with controllable memory.