Pipeline Parallelism and the Bubble

The problem

With pipeline parallelism, the layers of a model are divided among $p$ GPUs (graphics processing units), called stages: stage 0 holds the first layers, stage 1 the next ones, and so on. Each stage can only work on a piece of data after the previous stage has finished with it. To keep the stages busy, a batch is cut into $m$ micro-batches that follow each other through the stages, like cars on an assembly line. You will build GPipe's forward schedule and measure how much time the stages spend idle.

The schedule

Time is divided into slots, and every forward pass of one micro-batch on one stage takes one slot. Micro-batch $j$ reaches stage $s$ at slot $s + j$: stage 0 starts micro-batch 0 at slot 0, stage 1 can start it only at slot 1, after stage 0 is done with it.

The schedule is a table with one row per stage and one column per slot; it needs $m + p - 1$ slots for the last micro-batch to leave the last stage.

Example. $p = 3$ stages, $m = 2$ micro-batches: $2 + 3 - 1 = 4$ slots. -1 means idle.

stage	slot 0	slot 1	slot 2	slot 3
0	0	1	-1	-1
1	-1	0	1	-1
2	-1	-1	0	1

6 of the 12 cells are idle: the stages spend half of the time waiting.

The bubble

The idle cells are the bubble: at the start the pipeline fills (later stages wait), at the end it drains (earlier stages have nothing left). Each stage works for $m$ slots out of $m + p - 1$, so the idle fraction is:

where $p$ is the number of stages and $m$ the number of micro-batches.

Example. For the table above, $(3 - 1) / (2 + 3 - 1) = 2/4 = 1/2$, matching the 6 idle cells out of 12. With $p = 3$ and $m = 30$ it drops to $2 / 32 \approx 6\%$: the bubble shrinks only when there are many more micro-batches than stages.

What to implement

• gpipe_forward_schedule(p, m): a NumPy integer array of shape $(p, m + p - 1)$, where entry [s, t] is the micro-batch that stage $s$ processes in slot $t$, or -1 if the stage is idle.
• bubble_fraction(schedule): the fraction of all cells that are idle, measured from the schedule.
• bubble_formula(p, m): the formula above.

Requirements

• Stage $s$ starts micro-batch $j$ exactly one slot after stage $s - 1$.
• bubble_fraction counts the -1 cells; it must agree with bubble_formula for every $p$ and $m$.

Hint

Row $s$ of the table is $s$ idle slots, then the micro-batches $0, \dots, m - 1$, then idle slots until the end. Fill the array with -1, then write np.arange(m) into the right slice of each row.