colossalai.context.process_group_initializer.initializer_1d

class colossalai.context.process_group_initializer.initializer_1d.Initializer_1D(*args, **kwargs)

A ProcessGroupInitializer for 1d tensor parallelism.

init_dist_group()

Initialize 1D tensor parallel groups, and assign local_ranks and groups to each gpu. :return: (local_rank, group_world_size, process_group, ranks_in_group, mode) :rtype: Tuple