colossalai.context.process_group_initializer.initializer_model

class colossalai.context.process_group_initializer.initializer_model.Initializer_Model(*args, **kwargs)[source]

A ProcessGroupInitializer for model parallelism (model parallel group contains pipeline and tensor parallel groups).

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

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

Returns

A Model parallelism’s information tuple.

Return type

Tuple (local_rank, group_world_size, process_group, ranks_in_group, mode)