Schedule Pipeline Microbatches
Medium · Solved in PyTorch · Machine learning coding practice
Problem
Pipeline parallelism splits the layers across P stages and the global batch into M microbatches. The default Megatron schedule is non-interleaved one-forward-one-backward (1F1B), which keeps the same total bubble as filling the whole pipeline first but bounds how many activations a stage must store (Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM (Narayanan et al., 2021)).
Topics: megatron, pipeline-parallel, scheduling, pytorch, numpy
The full statement, worked examples, hints and the test suite are available once you sign in. You can then solve Schedule Pipeline Microbatches in the browser and run it against the tests.