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

• args – Args used to initialize ProcessGroupInitializer

• kwargs – Kwargs used to initialize ProcessGroupInitializer

init_dist_group()

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

Returns

Pipeline parallelism’s information

Return type

list of Tuples (local_rank, group_world_size, process_group, ranks_in_group, mode)