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

• args – Args used in base class

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

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

• args – Args used in base class

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

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

• args – Args used in base class

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

Return type

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

class colossalai.context.process_group_initializer.initializer_3d.Initializer_3D(*args)

Serve as the single entry point to 3D parallel initialization. :param args: Args used to initialize ProcessGroupInitializer

init_dist_group()

Initialize 3D tensor parallel groups, and assign local_ranks and groups to each gpu. :return: 3D tensor parallelism’s information :rtype: list of Tuples (local_rank, group_world_size, process_group, ranks_in_group, mode)