Memory-Aware Joint Optimisation of Partitioning and Scheduling for Pipeline-Parallel Training

Abstract

Pipeline parallelism (PP) is widely adopted for distributed training of neural networks, but its efficiency is limited by pipeline bubbles that leave devices idle. A variety of schedules have been proposed to reduce these bubbles, yet most are hand-crafted, rely on idealised timing assumptions as well as rigid structural constraints, and cannot adjust to memory limitations. Moreover, stage partitioning, task scheduling, and memory management are often optimised in isolation, despite being tightly interdependent. In this paper, we present OptPipe, a unified framework that jointly optimises partitioning and scheduling for pipeline parallelism. We introduce a memory-aware directed acyclic graph (DAG) that captures both task dependencies and the lifetime of intermediate tensors, enabling explicit reasoning about the trade-off between execution efficiency and memory usage. Building on this representation, we present a mathematical formulation for pipeline parallelism that simultaneously determines stage partitioning, device assignment, task ordering, and memory feasibility. Unlike prior methods, our approach imposes no structural constraints on pipeline configurations and adapts naturally to a wide range of memory budgets and workload characteristics. Experiments on GPT-3 and LLaMA-3.2 show that OptPipe improves training throughput by up to 72.5%, 37.8%, and 38.0% over 1F1B, Chimera, and ZB-V, while remaining feasible under memory budgets where these baselines fail.