colossalai.context.process_group_initializer.initializer_model

class colossalai.context.process_group_initializer.initializer_model.Initializer_Model(*args, **kwargs)

A ProcessGroupInitializer for model parallelism (model parallel group contains pipeline and tensor parallel groups).

Parameters

• args – Args used to initialize ProcessGroupInitializer

• kwargs – Kwargs used to initialize ProcessGroupInitializer

init_dist_group()

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

Returns

(local_rank, group_world_size, process_group, ranks_in_group, mode)

Return type

Tuple