Pipeline Warmup Micro-Batches

Problem

Count the warmup and steady-state micro-batches of a 1F1B pipeline schedule. Warmup fills the pipeline with min(num_stages - 1, num_microbatches) forward-only micro-batches; everything after warmup is steady state where forward and backward alternate. Return [warmup, steady].

The signature is pipeline_warmup_microbatches(num_stages, num_microbatches).

Hint

A shorter micro-batch stream can finish entirely inside the warmup phase.

Test cases

case 1

in: [4,8]

expected: [3,5]

case 2

in: [4,2]

expected: [2,0]

case 3

in: [1,5]

expected: [0,5]

case 4

in: [8,8]

expected: [7,1]

case 5

in: [2,5]

expected: [1,4]

Starter codePython

def pipeline_warmup_microbatches(num_stages, num_microbatches):
    # Your code here
    pass

5 test cases · run your code against them in the editor