Add 1F1B schedule

[kwen2501]

Description

c10d version 1F1B schedule.
We use nstages chunks to warm up the pipeline, then we enter a 1F1B stable phase, and lastly use nstages to cool down the pipeline.
Ref: https://arxiv.org/pdf/2104.04473.pdf.

API

stage = compile_stage(
    model,
    ... ,
    schedule="1F1B",
)

Implementation Details

To avoid writing duplicated code, we modularize the original code in FillDrain implementation into forward_one_chunk and backward_one_chunk. Then the two different schedules can share these modular code and just focus on the schedule itself.

Test

Added "schedule" option to test:

$ torchrun --nproc-per-node 4 local_test_c10d_bwd.py --schedule=1F1B --chunks=16