- class colossalai.pipeline.PipelinableContext(policy='balanced')
A context manager to split the model into pipeline stages.
Create a layer spec list and func list with execution sequence given by user. If exec_seq is None, we will take the module initizing order as execution order.
- partition(num_chunks, pipeline_size, rank)
Partitioned model will be built respect to partion policy. The real module instance will be built in this method.