colossalai.context.process_group_initializer.initializer_pipeline

class colossalai.context.process_group_initializer.initializer_pipeline.Initializer_Pipeline(*args, **kwargs)

A ProcessGroupInitializer for pipeline parallelism.

Parameters

• rank (int) – The rank of current process

• world_size (int) – Size of whole communication world

• config (Config) – Running configuration

• data_parallel_size (int) – Size of data parallel

• pipeline_parallel_size (int) – Size of pipeline parallel

• tensor_parallel_size (int) – Size of tensor parallel

init_dist_group()

Initialize pipeline parallel groups, and assign local_ranks and groups to each gpu.

Returns

A Pipeline parallelism’s information in list of tuples.

Return type

List[Tuple (local_rank, group_world_size, process_group, ranks_in_group, mode)]