class colossalai.pipeline.PipelinableContext(policy='balanced')[source]

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)[source]

Partitioned model will be built respect to partion policy. The real module instance will be built in this method.

class colossalai.pipeline.LayerSpec(typename, *module_args, **module_kwargs)[source]

Build the stored specification.