colossalai.context.process_group_initializer.initializer_1d

class colossalai.context.process_group_initializer.initializer_1d.Initializer_1D(*args, **kwargs)[source]

A ProcessGroupInitializer for 1d tensor 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()[source]

Initialize 1D tensor parallel groups, and assign local_ranks and groups to each gpu.

Returns

1D tensor parallelism’s information in a tuple.

Return type

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