colossalai.context.process_group_initializer.initializer_3d

class colossalai.context.process_group_initializer.initializer_3d.Initializer_3D_Input(num_group, depth, *args)

3D tensor parallel initialization among input.

Parameters

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

• depth (int) – Depth of 3D parallelism.

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

Returns

3D tensor parallelism’s information among input in a tuple.

Return type

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

class colossalai.context.process_group_initializer.initializer_3d.Initializer_3D_Weight(num_group, depth, *args)

3D tensor parallel initialization among weight.

Parameters

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

• depth (int) – Depth of 3D parallelism.

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

Returns

3D tensor parallelism’s information among weight in a tuple.

Return type

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

class colossalai.context.process_group_initializer.initializer_3d.Initializer_3D_Output(num_group, depth, *args)

3D tensor parallel initialization among output.

Parameters

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

• depth (int) – Depth of 3D parallelism.

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

Returns

3D tensor parallelism’s information among output in a tuple.

Return type

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