A pipeline-parallel training schedule planner and discrete-event simulator that picks the schedule with the least idle time under a fixed activation-memory budget.
When a model is too large for one accelerator, its layers are split across several devices as a pipeline. Microbatches flow through the stages so that every device has something to do. Two costs fight each other:
- Pipeline bubble: while the pipeline fills at the start and drains at the end, some devices sit idle. The bubble shrinks as you push more microbatches through per optimizer step.
- Activation memory: every forward pass whose backward has not run yet must keep its activations resident. The more microbatches are in flight, the more memory each device holds.
The classic GPipe schedule runs all forward passes and then all backward passes. That keeps the code simple, but the first stage has to hold the activations of every microbatch at once. Under a real memory limit you are forced to use very few microbatches per step, which means a large bubble and idle hardware. Picking the schedule and microbatch count by hand, per model and per cluster, is easy to get wrong, and the wrong choice quietly wastes a large fraction of the GPUs you are paying for.
pipebubble models a synchronous pipeline as a set of forward and backward tasks with data dependencies (forward flows from the first stage upward, backward flows from the last stage down) and simulates any schedule to completion. From one run it reports:
- makespan and per-device busy time
- pipeline bubble fraction and compute utilization
- peak activation memory, computed as the maximum overlap of live activation intervals per device
On top of the simulator it provides two schedule generators, GPipe and 1F1B (the PipeDream-Flush one-forward-one-backward schedule), and a planner. Given the pipeline depth, the number of microbatches, and a per-device activation-memory budget, the planner enumerates the feasible schedules and returns the one with the lowest bubble that stays within the budget.
The simulator is validated against the analytic pipeline-bubble formula and the analytic peak-memory bounds, so the numbers it produces are the same ones the theory predicts, not a hand-wave.
The insight the planner turns into a decision is that the two schedules scale differently with memory:
- GPipe peak memory equals the number of microbatches in a flush.
- 1F1B peak memory equals the pipeline depth, no matter how many microbatches stream through.
That asymmetry means that under a tight memory budget the two schedules are not close. GPipe is forced to chop the batch into many small flushes, each paying the full fill-and-drain bubble, while 1F1B streams the whole batch through a single fill-and-drain at the same memory. The planner makes that trade explicit and picks the winner for a given cluster instead of leaving it to a rule of thumb.
Architecture board on FigJam: https://www.figma.com/board/w7K7fapi7GTiUIxvnybEQl
The diagram shows how a config (stages, microbatches, memory budget, per-pass costs) flows through the schedule generators into the discrete-event simulator, how the simulator's dependency resolver and per-device resource constraint produce the timing and memory metrics, and how the planner selects the lowest-bubble feasible schedule. A separate branch shows the validation check that ties the simulated bubble back to the closed-form result.
graph TD
A["config: P, N, memory budget"] --> B["schedule generator: GPipe or 1F1B"]
B --> C["discrete-event simulator"]
C --> D["timing: makespan, bubble, utilization"]
C --> E["memory: peak live activations per device"]
D --> F["planner: lowest bubble within budget"]
E --> F
F --> G["chosen schedule and microbatch grouping"]
The metric is pipeline compute utilization: the fraction of device time spent doing forward and backward work rather than idling in the bubble. The setup is an 8 stage pipeline processing a global batch of 64 microbatches, with each device limited to 8 activation slots and a backward pass costing twice a forward pass. Both schedules are held to the same 8 slot memory budget, so this is not a memory-for-speed trade: the memory panel shows both peaking at exactly 8 slots.
Under that budget GPipe can hold at most 8 microbatches, so it processes the 64 microbatch batch as 8 separate fill-and-drain flushes and pays the bubble eight times, landing at 53.3 percent utilization (46.7 percent idle). 1F1B holds at most 8 activations regardless of how many microbatches flow through, so it streams all 64 in a single flush and pays the bubble once, reaching 90.1 percent utilization (9.9 percent idle). That is a 1.69x wall-clock speedup at identical memory. The planner selects 1F1B automatically.
The Gantt view below (regenerated by examples/plot_gantt.py) shows why, on a smaller 4 stage, 8 microbatch case. GPipe holds 8 activations at the first stage; 1F1B holds 4 (the pipeline depth) at the same bubble.
Install the one dependency (only the plotting scripts need it; the simulator and planner are pure standard library):
pip install -r requirements.txt
Run the validation tests, which check the simulator against the closed-form bubble formula and the analytic memory bounds:
python tests/test_simulator.py
Run the benchmark, which prints the numbers above and regenerates docs/before_after.png:
python examples/run_benchmark.py
Regenerate the Gantt figure:
python examples/plot_gantt.py
Use the planner directly:
from pipebubble import plan
result = plan(num_stages=8, total_microbatches=64, memory_budget=8)
print(result["chosen"]["schedule"], result["chosen"]["utilization"])Pipeline utilization under a fixed activation-memory budget
--------------------------------------------------------------
stages (P) : 8
microbatches (N) : 64
memory budget per device: 8 activation slots
forward / backward cost : 1.0 / 2.0
GPipe (baseline)
flushes : 8
peak activation mem: 8 slots
makespan : 360.0
pipeline bubble : 46.7 %
utilization : 53.3 %
1F1B (planner)
flushes : 1
peak activation mem: 8 slots
makespan : 213.0
pipeline bubble : 9.9 %
utilization : 90.1 %
planner chose : 1f1b
utilization : 53.3 % -> 90.1 %
bubble : 46.7 % -> 9.9 %
wall-clock speedup : 1.69x at equal memory budget