Overview

1. We add interleaved pipeline schedule.
2. Update pipeline model initializer.
3. Fix Naive AMP

Interleaved pipeline schedule and pipeline model initializer

We kept the original PipelineModelInitializer, but renamed it to build_pipeline_model_from_cfg. We add a new function to build the pipeline parallel model: build_pipeline_model. User must pass a torch.nn.Sequential to this function.

Usage:

from colossalai.engine.schedule import InterleavedPipelineSchedule
from colossalai.builder import build_pipeline_model_from_cfg, build_pipeline_model

# from config
model = build_pipeline_model_from_cfg(model_cfg, num_model_chunks)

# from torch.nn.Sequential
model = build_pipeline_model(sequential_model, num_model_chunks)

schedule = InterleavedPipelineSchedule(num_microbatches, num_model_chunks)

build_pipeline_model and build_pipeline_model_from_cfg can also receive a key word argument verbose. The detail of layers in each rank will be printed if setting verbose to True.

Fix Naive AMP

The old Naive AMP didn't apply all-reduce over pipeline parallel process group when computing infinity values, which leads to a hang. I fixed this bug, so that both pipeline parallelism and interleaved pipeline parallelism with Naive AMP won't get stuck again. Note that the pipeline parallelism only support Naive AMP now.