colossalai.pipeline

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

A context manager to split the model into pipeline stages.

to_layer_list(exec_seq=None)[source]: 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]