colossalai.context.process_group_initializer.initializer_tensor

class colossalai.context.process_group_initializer.initializer_tensor.Initializer_Tensor(*args, **kwargs)

A ProcessGroupInitializer for tensor parallelism.

Parameters

• args – Args used to initialize ProcessGroupInitializer

• kwargs – Kwargs used to initialize ProcessGroupInitializer

init_dist_group()

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

Returns

Tensor parallelism’s information

Return type

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