colossalai.context.process_group_initializer.initializer_tensor

class colossalai.context.process_group_initializer.initializer_tensor.Initializer_Tensor(*args, **kwargs)[source]

A ProcessGroupInitializer for 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 tensor parallel groups, and assign local_ranks and groups to each gpu.

Returns

A Tensor parallelism’s information tuple.

Return type

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