专家混合模型
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 运行时中,并默认启用,因此用户无需任何配置标志或命令行参数即可使用。