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. :param num_group: The number of all tensor groups :param summa_dim: The dimension of SUMMA :param args: Args used to initialize base class :param kwargs: Kwargs used to initialize base class :type num_group: int :type summa_dim: int
- 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
- Return type
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
args – Args used to initialize base class
kwargs – Kwargs used to initialize base class
- 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
- Return type
Tuple(local_rank, group_world_size, process_group, ranks_in_group, mode)
- class colossalai.context.process_group_initializer.initializer_2d.Initializer_2D(*args, **kwargs)
Serve as the single entry point to 2D parallel initialization.
- Parameters
args – Args used to initialize ProcessGroupInitializer
kwargs – Kwargs used to initialize ProcessGroupInitializer
- init_dist_group()
Initialize 2D tensor row and col parallel groups, and assign local_ranks and groups to each gpu. :return: 2D tensor parallelism’s information :rtype: list of Tuples (local_rank, group_world_size, process_group, ranks_in_group, mode)