colossalai.context.moe_context

class colossalai.context.moe_context.MoeParallelInfo(ep_size, dp_size)

Moe parallelism information, storing parallel sizes and groups.

class colossalai.context.moe_context.MoeContext(*args, **kwargs)

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

get_info(num_experts)

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