专家混合模型

DeepSpeed v0.5 引入了对训练专家混合模型 (MoE) 的新支持。MoE 模型是一类新兴的稀疏激活模型，其计算成本相对于其参数而言是次线性的。例如，Switch Transformer 包含超过 1.6 万亿个参数，而训练它所需的计算量大约等于一个 100 亿参数的密集模型。模型规模的这种增加为恒定的计算预算提供了巨大的精度提升。

有关结果和进一步讨论的更多详细信息，请参阅我们的新闻稿：DeepSpeed 实现了 8 倍更大的 MoE 模型训练性能。

使用简单的 MoE 示例入门

注意：DeepSpeed MoE 需要 Pytorch 1.8 或更高版本。

作为一个简单的起点，我们将展示如何将 DeepSpeed MoE 应用于 cifar10 示例。请参考我们的 cifar10 示例 以便后续操作。

如果您正在将 MoE 添加到现有模型中，可以使用下面的代码段来指导您。

专家组初始化

DeepSpeed MoE 支持五种不同的并行形式，并利用 GPU 和 CPU 内存。其灵活的设计使用户能够混合使用不同类型的流行并行技术，如下表所示。

简称 灵活的并行配置 优势
E 专家 通过增加专家数量来扩展模型规模
E + D 专家 + 数据 通过扩展到多个数据并行组来加速训练吞吐量
E + Z 专家 + ZeRO 支持的数据 对非专家参数进行分区以支持更大的基础模型
E + D + M 专家 + 数据 + 模型 支持更大的隐藏大小，以及比 E+Z 更大的基础模型
E + D + Z 专家 + 数据 + ZeRO 支持的数据 支持更大的隐藏大小，以及比 E+Z 更大的基础模型
E + Z-Off + M 专家 + ZeRO-Offload + 模型 利用 GPU 和 CPU 内存，在有限数量的 GPU 上训练大型 MoE 模型

为了支持不同形式的并行性，我们在 DeepSpeed 内部创建了各种进程组。DeepSpeed 使用的辅助函数位于 deepspeed.utils.groups.py 中

注意：以下函数现已弃用，模型训练代码不再需要调用它。

deepspeed.utils.groups.initialize(ep_size="desired expert-parallel world size")

相反，MoE 层 API 现在除了 num_experts 之外还接受 ep_size 作为参数。这个新的 API 允许用户创建 MoE 模型，这些模型可以具有不同数量的专家，并且每个 MoE 层可以具有不同的专家并行度。

参与大小为 ep_size 的专家并行组的 GPU（或排名）将分配层指定的专家总数。

MoE 层 API

hidden_size 是特定层的输入维度，输出维度与之相同。这可能导致模型定义的一些变化，特别是对于视觉/卷积模型，因为在某些情况下输入/输出维度不匹配。例如，在 CIFAR-10 示例中，我们修改了第三个全连接层以添加 MoE 层。为了满足这一需求，我们需要添加一个额外的全连接层，其输入维度等于 MoE 层的输出维度。

原始模型配置

    self.fc3 = nn.Linear(84, 10)

使用 MoE 层更新

    self.fc3 = nn.Linear(84, 84)
    self.fc3 = deepspeed.moe.layer.MoE(hidden_size=84, expert=self.fc3, num_experts=args.num_experts, ep_size=<desired expert-parallel world size> ...)
    self.fc4 = nn.Linear(84, 10)

金字塔残差 MoE

最近，我们提出了一种新颖的 金字塔残差 MoE (PR-MoE) 模型架构。要创建这样的 MoE 模型，用户需要做两件事：1) 要创建金字塔结构，请将 num_experts 作为列表传递，例如 [4, 8]；2) 使用 use_residual 标志指示 MoE 层现在是残差 MoE 层。

self.experts = deepspeed.moe.layer.MoE(hidden_size=input_dim, expert=ExpertModule(), num_experts=[..], ep_size=ep_size, use_residual=True)

示例场景

给定世界大小中的 GPU 总数以及专家并行世界中的 GPU 子集，如下所示。

WORLD_SIZE = 4
EP_WORLD_SIZE = 2
EXPERTS = [8]

模型代码需要使用 deepspeed.moe.layer.MoE API，如下所示。

self.experts = deepspeed.moe.layer.MoE(hidden_size=input_dim, expert=ExpertModule(), num_experts=EXPERTS, ep_size=EP_WORLD_SIZE)

通过以上两个命令，DeepSpeed 运行时将设置为训练一个 MoE 模型，该模型在 4 个 GPU 上总共有 8 个专家，以 4 个专家/GPU 的模式运行。我们将其称为 E + D 模式，如上表所述。

import torch
import deepspeed
import deepspeed.utils.groups as groups
from deepspeed.moe.layer import MoE

WORLD_SIZE = 4
EP_WORLD_SIZE = 2
EXPERTS = 8

fc3 = torch.nn.Linear(84, 84)
fc3 = MoE(hidden_size=84, expert=self.fc3, num_experts=EXPERTS, ep_size=EP_WORLD_SIZE, k=1)
fc4 = torch.nn.Linear(84, 10)

有关涵盖标准 MoE 架构和 PR-MoE 模型的端到端可运行示例，请查看 cifar10 示例。此外，请参阅本教程的高级用法部分，该部分链接到用于 NLG 模型的更全面的示例。

结合 ZeRO-Offload 和 DeepSpeed MoE 用于超大型模型

要在 DeepSpeed 中使用 MoE 层，我们依赖于传递给优化器的两个参数组。从 cifar10 示例 中可以找到创建此类组的具体示例。

创建这些参数组的相关函数如下所示。

def create_moe_param_groups(model):
    from deepspeed.moe.utils import split_params_into_different_moe_groups_for_optimizer

    parameters = {'params': [p for p in model.parameters()], 'name': 'parameters'}

    return split_params_into_different_moe_groups_for_optimizer(parameters)

然后，上述参数组可以馈送到 ZeRO 第 2 阶段优化器，如下所示。

net = Net()

parameters = create_moe_param_groups(net)

model_engine, optimizer, trainloader, __ = deepspeed.initialize(
    args=args, model=net, model_parameters=parameters, training_data=trainset)

我们正在努力在 DeepSpeed ZeRO 优化器中自动化此功能，以便进一步简化模型训练代码。

要使用 ZeRO-Offload（第 2 阶段）和 MoE 运行 cifar10 示例，请设置 ds_config 标志

"zero_optimization": {
      "stage": 2,
      "allgather_partitions": true,
      "reduce_scatter": true,
      "allgather_bucket_size": 50000000,
      "reduce_bucket_size": 50000000,
      "overlap_comm": true,
      "contiguous_gradients": true,
      "cpu_offload": true
  }

还引入了一种额外的优化方法，以节省在有限数量的 GPU 上训练超大型模型的内存。请使用以下配置标志将其启用到 ds_config 中的 fp16 优化器。

    "fp16": {
      "enabled": true,
      "fp16_master_weights_and_grads": true,
  }

随机标记选择

我们设计了一种名为“随机标记选择”的新技术，可以极大地提高收敛速度。随机标记选择解决了 MoE 模型训练中存在的有偏选择问题的限制。我们即将发表的论文将详细介绍此技术及其结果。此功能已成为 DeepSpeed 运行时的一部分，并且默认启用，因此用户无需任何配置标志或命令行参数即可利用此功能。

高级 MoE 用法

我们添加了将 MoE 应用于 NLG 模型的示例。请在此 新闻稿 和 教程 中了解更多信息。

更新：