colossalai.builder.pipeline

colossalai.builder.pipeline.count_layer_params(layers)

Count the number of parameters in each layer

colossalai.builder.pipeline.build_pipeline_model_from_cfg(config, num_chunks=1, partition_method='parameter', verbose=False)

An intializer to split the model into different stages for pipeline parallelism.

An example for the model config is shown below. The class VisionTransformerFromConfig should inherit colossalai.nn.model.ModelFromConfig to allow this initializer to build model from a sequence of layer configurations.

model_config = dict(

type=’VisionTransformerFromConfig’, embedding_cfg=dict(…), …

)

Parameters

• config (dict) – Configuration of the model

• num_chunks (int, optional) – The number of chunks you want to have on the current stage. This value should be 1 in most cases unless you are using virutal pipeline parallelism.

• partition_method (str, optional) – This parameter determines how you want to split your model layers into stages, you can set it as ‘layer’ or ‘parameter’

• verbose (bool, optional) – Whether to print the logs

colossalai.builder.pipeline.build_pipeline_model(layers, num_chunks=1, verbose=False)

An intializer to split the model into different stages for pipeline parallelism. Note that layer must be torch.nn.Sequential.

Parameters

• layers (torch.nn.Sequential) – Layers of model

• num_chunks (int, optional) – The number of chunks you want to have on the current stage. This value should be 1 in most cases unless you are using virutal pipeline parallelism.

• verbose (bool, optional) – Whether to print the logs