OptPipe: Memory- and Scheduling-Optimized Pipeline Parallelism for LLM Training

TL;DR

OptPipe tackles memory and scheduling bottlenecks in pipeline parallelism for LLM training by formulating the problem as a MILP that jointly optimizes computation timing, activation offloading, and memory usage under device constraints. It introduces AdaOffload, cached schedules, and online scheduling with topology-aware constraints to efficiently navigate the memory-time trade-offs. The approach yields substantial throughput improvements and enables training larger models within fixed memory budgets, including up to a 50% reduction in idle pipeline time under the same memory cap. Empirical results on Megatron-LM-like models across up to 16 H100 GPUs demonstrate robust gains over baselines, highlighting the practical impact of refined, MILP-based scheduling for real-world large-scale training.

Abstract

Pipeline parallelism (PP) has become a standard technique for scaling large language model (LLM) training across multiple devices. However, despite recent progress in reducing memory consumption through activation offloading, existing approaches remain largely heuristic and coarse-grained, often overlooking the fine-grained trade-offs between memory, computation, and scheduling latency. In this work, we revisit the pipeline scheduling problem from a principled optimization perspective. We observe that prevailing strategies either rely on static rules or aggressively offload activations without fully leveraging the interaction between memory constraints and scheduling efficiency. To address this, we formulate scheduling as a constrained optimization problem that jointly accounts for memory capacity, activation reuse, and pipeline bubble minimization. Solving this model yields fine-grained schedules that reduce pipeline bubbles while adhering to strict memory budgets. Our approach complements existing offloading techniques: whereas prior approaches trade memory for time in a fixed pattern, we dynamically optimize the tradeoff with respect to model structure and hardware configuration. Experimental results demonstrate that our method consistently improves both throughput and memory utilization. In particular, we reduce idle pipeline time by up to 50% under the same per-device memory limit, and in some cases, enable the training of larger models within limited memory budgets.

Figures (6)

• Figure 1: The framework of OptPipe. The framework consists of three main phases: (1) Initialize: Generate an initial scheduling strategy that ensures peak memory usage remains within device limits; (2) Profile: Run a few warm-up iterations to profile computation time and memory usage, and use the collected statistics to construct the MILP model; (3) Schedule & Train: Implement an initial schedule using a heuristic strategy (e.g., PipeOffload or a cached schedule), and employ an MILP solver to refine the schedule. Training proceeds in parallel, with updates applied whenever the solver discovers an improved solution.
• Figure 2: The computation graph of each stage in pipeline parallelism.
• Figure 3: Determined and Undetermined Variables
• Figure 4: A toy example for illustration. Each block in the figure represents a distinct operation type, with the colors indicating different stages. The x-axis represents time, and the y-axis shows different stages in the pipeline parallelism.
• Figure 5: Memory Usage Comparison between PipeOffload and OptPipe. Average, minimum, and maximum device memory usage across different model sizes and micro-batch sizes.