专家混合模型 (MoE)

图层规范

class deepspeed.moe.layer.MoE(hidden_size: int, expert: Module, num_experts: int = 1, ep_size: int = 1, k: int = 1, capacity_factor: float = 1.0, eval_capacity_factor: float = 1.0, min_capacity: int = 4, use_residual: bool = False, noisy_gate_policy: Optional[str] = None, drop_tokens: bool = True, use_rts: bool = True, use_tutel: bool = False, enable_expert_tensor_parallelism: bool = False, top2_2nd_expert_sampling: bool = True)

初始化一个MoE层。

Parameters

• hidden_size (int) – 模型的隐藏维度，重要的是这也是输入和输出的维度。

• 专家 (nn.Module) – 定义专家的torch模块（例如，MLP，torch.linear）。

• num_experts (int, optional) – 默认=1，每层的专家总数。

• ep_size (int, optional) – 默认=1，专家并行世界或组中的排名数量。

• k (int, 可选) – 默认=1, top-k 门控值，仅支持 k=1 或 k=2。

• capacity_factor (float, optional) – 默认值=1.0，训练时专家的容量。

• eval_capacity_factor (float, optional) – 默认值=1.0，评估时专家的能力。

• min_capacity (int, optional) – 默认=4，无论capacity_factor如何，每个专家的最小容量。

• use_residual (bool, 可选) – 默认=False, 使这个MoE层成为一个残差MoE (https://arxiv.org/abs/2201.05596) 层。

• noisy_gate_policy (str, optional) – 默认=None, 噪声门策略，有效选项为‘Jitter’, ‘RSample’ 或 ‘None’。

• drop_tokens (bool, optional) – 默认=True, 是否丢弃令牌 - (设置为False相当于无限容量).

• use_rts (bool, optional) – 默认=True, 是否使用随机令牌选择。

• use_tutel (bool, optional) – 默认=False，是否使用Tutel优化（如果已安装）。

• enable_expert_tensor_parallelism (bool, optional) – 默认=False，是否对专家使用张量并行

• top2_2nd_expert_sampling (bool, optional) – 默认=True, 是否对第二个专家进行采样

forward(hidden_states: Tensor, used_token: Optional[Tensor] = None) → Tuple[Tensor, Tensor, Tensor]

MoE 前向传播

Parameters

• hidden_states (Tensor) – 输入到该层的内容

• used_token (Tensor, optional) – 默认值: None, 仅屏蔽已使用的令牌

Returns

一个包含输出、门损失和专家计数的元组。

• output (Tensor): 模型的输出

• l_aux (Tensor): 门损失值

• exp_counts (Tensor): 专家计数