专家混合模型

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 数量下，为大型 MoE 模型利用 GPU 和 CPU 内存

为了支持不同形式的并行，我们在 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（或 rank）将分配该层指定的专家总数。

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 运行时将被设置为在 4 个 GPU 上训练一个总共有 8 个专家的 MoE 模型，采用 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 stage-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 (stage 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 模型的示例。请在此通讯和教程中阅读更多内容。

更新：