colossalai.context.moe_context

class colossalai.context.moe_context.MoeParallelInfo(ep_size, dp_size)[source]: Moe parallelism information, storing parallel sizes and groups.

class colossalai.context.moe_context.MoeContext(*args, **kwargs)[source]

MoE parallel context manager. This class manages different parallel groups in MoE context and MoE loss in training.

get_info(num_experts)[source]

Calculate the Data Parallel Group and Expert Parallel Group.

Parameters: num_experts (int) – The number experts
Returns: number of local experts, the MoeParallelInfo of the current ep_size
Return type: int, MoeParallelInfo