重要提示

您正在查看 NeMo 2.0 文档。此版本对 API 和一个新库 NeMo Run 进行了重大更改。我们目前正在将 NeMo 1.0 的所有功能移植到 2.0。有关先前版本或 2.0 中尚不可用的功能的文档，请参阅 NeMo 24.07 文档。

专家混合模型#

概述#

NeMo 框架支持 transformer 层的 feedforward 块中的专家混合模型 (MoE)。

MoE 是一种机器学习技术，其中结合了多个专门的模型（专家，通常是多层感知器）来解决复杂的任务。每个专家专注于特定的子任务或领域，而门控网络根据当前的输入动态激活最合适的专家。

使用 MoE#

要在 NeMo 框架中使用 MoE，请调整模型配置中的 num_moe_experts 参数

  1. 将 num_moe_experts 设置为 8 以在 MoE 模块中利用 8 个专家。

    num_moe_experts: 8  # Set MoE to use 8 experts

  2. 将 moe_router_topk 设置为您想要激活的专家数量。例如，如果您想使用两个专家处理每个输入

    moe_router_topk: 2  # Processes each token using 2 experts.

配置 MoE 特定的损失函数#

此外，NeMo 提供了配置 MoE 特定的损失函数的选项。为了平衡专家之间的令牌分布

  1. 设置 moe_router_load_balancing_type 以指定负载均衡方法

    moe_router_load_balancing_type: aux_loss  # to use the auxilary loss, other options include "sinkhorn".

  2. 设置 moe_aux_loss_coeff 以指定辅助损失的权重。添加辅助损失是为了鼓励在所有专家之间平均分配令牌。1e-2 范围内的值是一个好的开始，如下所示

    moe_aux_loss_coeff: 1e-2  # set the aux-loss weight to 1e-2

  3. 设置 moe_z_loss_coeff 以指定 z-loss 的权重。建议的起始值为 1e-3，如下所示

    moe_z_loss_coeff: 1e-3

其他选项包括

  1. moe_input_jitter_eps 通过应用具有指定 epsilon 值的抖动，向输入张量添加噪声。

  2. moe_token_dropping 允许选择性地删除和填充每个专家的令牌，以实现指定的容量，类似于 GShard、Switch-Transformer 和 DeepSpeed-MoE。简而言之，如果路由到某个专家的令牌数量超过其容量，则会删除超出容量的令牌。请注意，目前不支持此功能，因此应保持为 False。

  3. moe_token_dispatcher_type 指定令牌调度器类型，选项包括“allgather”和“alltoall”。

  4. moe_per_layer_logging 启用 MoE 的逐层日志记录，目前支持 aux-loss 和 z-loss。

  5. moe_expert_capacity_factor 容量因子决定了在任何 MoE 层中可以路由到每个专家的最大令牌数量。None 表示不会删除任何令牌。默认值为 None。

  6. moe_pad_expert_input_to_capacity 如果为 True，则填充每个专家的输入以匹配专家容量长度。仅在设置 moe_expert_capacity_factor 后才有效。默认设置为 False。

  7. moe_token_drop_policy 删除令牌的策略。可以是“probs”或“position”。如果为“probs”，则将删除概率最低的令牌。如果为“position”，则将删除每个批次末尾的令牌。默认值为“probs”。

  8. moe_layer_recompute 如果为 True，则检查点 moe_layer 以节省激活内存。默认值为 False。

本页内容