Pipeline Parallelism with Controllable Memory

Summary

The paper addresses the inefficiencies in current pipeline parallelism schedules used for training large-scale models, particularly focusing on high activation memory and pipeline bubbles. The paper proposes a new framework to decompose pipeline schedules into repeating building blocks and introduces V-shape building blocks (V-Min, V-Half, V-ZB) with controllable activation memory. Such building blocks reduce peak activation memory compared to the traditional 1F1B approach and even achieve higher throughput with reduced pipeline bubbles. The evaluations demonstrate substantial improvements in throughput and memory efficiency over 1F1B and 1F1B-I methods.

Strengths

- The authors decompose pipeline parallelism into repeating building blocks, providing a systematic methodology for designing more efficient schedules. This approach helps in understanding and optimizing pipeline parallelism in a structured manner. - The V-Shape building blocks demonstrate the structured understanding of pipeline parallelism. - The authors provide experimental evaluations demonstrating the practical benefits of the proposed methods. - By reducing pipeline bubbles, the paper demonstrates that PP can become a much more preferred option over TP at practical large-scale training.

Weaknesses

- It would be helpful to have experimental results for long sequence lengths. 1024 and 3072 sequence lengths are too short compared to what SOTA LLMs can handle.

Questions

- One of the problems with pipeline parallelism is that it increases implementation complexity because of its multi-program-multi-data (MPMD) nature. Do V-shaped abstractions bring any convenience in terms of implementation perspective?

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

No societal impact exists.

Summary

This work proposes systematic methodology for designing pipeline parallelism schedules and analyzing their performance (e.g. peak memory and pipeline bubbles). The major observation is that a pipeline schedule can be viewed as repeating a building block, and the peak memory critically depends on the lifespan of a block. Based on these insights, the authors design a family of novel V-shape building blocks and pipeline schedules that are memory efficient and achieve higher throughput than strong baselines like 1F1B (though at the cost of more communication), which is validated both theoretically and empirically.

Strengths

- This work proposes a systematic and unified perspective for pipeline parallelism schedule, with some novel insights that might foster the design of better schedules in the future. - Analysis and experiments are extensive and well executed, confirming the efficacy of the proposed methodology. - Overall, writing and presentation are clear.

Weaknesses

There is still room for improvement in writing. For example: - There are quite some typos, some of which are listed here: Line 22 "tensor" -> "Tensor"; Line 189, "V-Min" -> "V-ZB" (?); Line 193, "serious" -> "series". - Line 210, "while other methods' memory is similar": one exception is 1F1B-R, whose activation memory is much smaller.

Questions

A high-level question that might be worth a brief discussion in the manuscript is: is it true that a pipeline schedule *must be* repetition of a building block? Actually the answer is no. A naive example I can think of is a variant of GPipe, where the order of backward passes can be arbitrary, and thus the building block for one microbatch can be different from that of another microbatch. This extension of GPipe is designed just out of theoretical curiosity though. So the next question is whether there could be any real advantage (in terms of peak memory, throughput, etc.) in a pipeline schedule that cannot be expressed as repetition of a building block. It would be great if there is a clear answer with a formal proof. But even if the authors do not know the answer yet (and neither do I), it might be worth mentioning, so that anyone who has read this work will keep this in mind, rather than blindly constrain themselves to repeating a building block when designing new schedules.

Rating

7

Confidence

3

Soundness

3

Presentation

2

Contribution

3

Limitations

Limitations are discussed throughout the manuscript.

Summary

Authors propose a way to identify the repeated building block a pipeline schedule is built from. By relating the peak activation memory of the schedule to the lifespan of the building block the authors show that existing schedules do not optimally use memory, and design higher-throughput schedules that use the same activation memory as the baseline. The methods demonstrate up to 55% improvement in throughput for pure pipeline parallelism settings and 16% improvement for large language models over the 1F1B baseline. The paper also provides a systematic methodology for designing pipeline schedules, emphasizing memory usage balance and reduced pipeline bubbles.

Strengths

The paper introduces a novel framework for decomposing pipeline parallelism schedules into memory-efficient building blocks. This approach addresses inefficiencies observed in existing schedules. The systematic methodology for designing pipeline schedules emphasizes memory usage balance and reduced pipeline bubbles. Discussion of peak memory and new memory efficient building blocks is the core of the paper and the part which is most clear. The authors provide technical details, including the proposed building blocks, their implementation, and experimental results.

Weaknesses

The paper suffers from quite a few clarity issues and could use some copyediting (there are lots of small grammatical errors, etc.). Main clarity issues in the paper are, for example, a lack of useful captions in figures, a lack of substantive discussion in some of the appendices (i.e., “we leave the discussion to appendix X” but there is not much discussion in appendix X). A lot could be done to improve clarity of the discussion in section 3. For example, although the paper explicitly details asymptotic behavior (d -> infinity?) this is really only mentioned in lines 166-167, which somewhat confuses the issue. Some brief discussion of what the effect is in low-d situations would be nice (not new experiments – just a qualitative idea). Fig. 3 and Table 1 contradict each other due to the presence/absence of the asymptotic limit, which is mentioned in the title of Table 1 and briefly in the text but is not very clear, especially since Fig. 3 and Table 1 are supposed to be read together (?). The figure captions should at minimum restate what is shown in the figures and what relevant terms (d, l, etc.) mean – this makes it much easier to refer to the figure without searching through the text for definitions and explanations. This should be doable in 1 or 2 sentences per caption at most and should not take a lot of space. Some plots are hard to read. For example, Figure 4 and Figure 5 shows detailed pipelines using various colors of blocks and fonts without proper definition or explanation. It is also not clear the shown pipeline is the actual setting or a high-level demonstration of the design. V-Half seems to be a heuristic method based on V-Min and V-ZB. It may not be an optimal solution for the pipeline. Other configurations regarding the trade-off of memory and throughput are not evaluated. Section 4.4 is confusing. Table 3 is not referenced anywhere and not explained in this section. The results mentioned in this section (Table 6) do not clearly show the combined performance of the proposed approach and existing techniques.

Questions

Why is 1F1B the baseline used for the paper? It feels like there could be some more discussion here that is missing. Is the lack of bubble in V-Half (figure 5) purely empirical or is there some reason it should be expected? There doesn’t seem to be anything special about V-Half from the point of view of the discussion in the paper (only the activation memory is relevant to V-half; it’s not clear why this specifically should not have bubbles). Also, V-ZB has no bubbles in the case of differing runtimes, if I understand correctly; so maybe it is more accurate to say that the number of bubbles decreases to 0 as activation memory budget increases? How is the number of stages in the pipeline determined? While the authors mentioned that it depends on the number of devices, what is the number of stages used in this work? How does the number of stages influence performance? Given the constraints of memory or throughput, do the authors use grid search to find an optimal schedule? Is it possible to use a theoretical or approximation method to quickly find a good pipeline schedule?

Rating

5

Confidence

3

Soundness

3

Presentation

2

Contribution

3

Limitations

Yes