colossalai.context.process_group_initializer.initializer_2p5d

class colossalai.context.process_group_initializer.initializer_2p5d.Initializer_2p5D_ROW(tesseract_dim, tesseract_dep, *args)

2p5d tensor parallel initialization among rows.

Parameters

• tesseract_dim (int) – The dimension of tesseract

• tesseract_dep (int) – The dimension of depth

• args – Args used to initialize base class

init_dist_group()

Initialize 2p5D tensor row parallel groups, and assign local_ranks and groups to each gpu.

Returns

2p5D tensor row parallelism’s information

Return type

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

class colossalai.context.process_group_initializer.initializer_2p5d.Initializer_2p5D_Col(tesseract_dim, tesseract_dep, *args)

2p5d tensor parallel initialization among cols.

Parameters

• tesseract_dim (int) – The dimension of tesseract

• tesseract_dep (int) – The dimension of depth

• args – Args used to initialize base class

init_dist_group()

Initialize 2p5D tensor col parallel groups, and assign local_ranks and groups to each gpu.

Returns

2p5D tensor col parallelism’s information

Return type

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

class colossalai.context.process_group_initializer.initializer_2p5d.Initializer_2p5D_Dep(tesseract_dim, tesseract_dep, *args)

2p5D tensor parallel initialization among depths.

Parameters

• tesseract_dim (int) – The dimension of tesseract

• tesseract_dep (int) – The dimension of depth

• args – Args used to initialize base class

init_dist_group()

Initialize 2p5D tensor depth parallel groups, and assign local_ranks and groups to each gpu.

Returns

2p5D tensor depth parallelism’s information

Return type

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

class colossalai.context.process_group_initializer.initializer_2p5d.Initializer_2p5D_XZ(tesseract_dim, tesseract_dep, *args)

2p5d tensor parallel initialization among cols times dep.

Parameters

• tesseract_dim (int) – The dimension of tesseract

• tesseract_dep (int) – The dimension of depth

• args – Args used to initialize base class

init_dist_group()

Initialize 2p5D tensor colXdepth parallel groups, and assign local_ranks and groups to each gpu.

Returns

2p5D tensor colXdepth parallelism’s information

Return type

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

class colossalai.context.process_group_initializer.initializer_2p5d.Initializer_2p5D(rank, world_size, config, data_parallel_size, pipeline_parallel_size, tensor_parallel_size, depth)

Serve as the single entry point to Tesseract parallel initialization.

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

• depth (int) – The depth of 2p5d parallel

init_dist_group()

Initialize 2p5D tensor row, col, depth, and colXdepth parallel groups, and assign local_ranks and groups to each gpu. :return: Whole 2p5D tensor parallelism’s information :rtype: list of Tuples (local_rank, group_world_size, process_group, ranks_in_group, mode)