colossalai.context.process_group_initializer.initializer_2d

class colossalai.context.process_group_initializer.initializer_2d.Initializer_2D_Row(num_group, summa_dim, *args, **kwargs)

2d tensor parallel initialization among rows.

Parameters

• num_group (int) – The number of all tensor groups.

• summa_dim (int) – The dimension of SUMMA.

• 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 2D tensor row parallel groups, and assign local_ranks and groups to each gpu. :returns: 2D tensor row parallelism’s information in a tuple. :rtype: Tuple (local_rank, group_world_size, process_group, ranks_in_group, mode)

class colossalai.context.process_group_initializer.initializer_2d.Initializer_2D_Col(num_group, summa_dim, *args, **kwargs)

2d tensor parallel initialization among cols.

Parameters

• num_group (int) – The number of all tensor groups.

• summa_dim (int) – The dimension of SUMMA.

• 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 2D tensor row parallel groups, and assign local_ranks and groups to each gpu.

Returns

2D tensor col parallelism’s information in a tuple.

Return type

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