User prompt:
You have access to source code for torch.distributed.pipelining
package and relevant documentation for implementing pipeline parallelism. Write a function that creates an accurate pipeline_order
for Schedule1F1B
given world_size
and n_microbatches
such that _format_pipeline_order
will return the correct output. Make sure to provide a minimal example to verify function performs as expected.
CoT:
- Crafting a pipeline: I’m thinking through a function to set up the pipeline order for Schedule1F1B in PyTorch, aiming to efficiently manage forward and backward passes across stages. Progress is steady, ensuring an interleaved pattern that optimizes pipeline stage utilization.
- Crafting the schedule: I’m working on a 1F1B pipeline schedule, alternating forward and backward pass microbatches at each stage. Ensuring optimal utilization, I aim to refine a formula for this arrangement.
- Piecing together: I’m using a well-known formula from “GPipe” and “torch.distributed.pipeline”