DeepSpeed-MoE 大规模 MoE 模型推理入门

DeepSpeed-MoE 推理在针对密集模型（DeepSpeed 推理博客文章）的基础上引入了多项重要功能。它涵盖了多种不同类型的并行，即针对非专家参数的数据并行和张量切片，以及针对专家参数的专家并行和专家切片。为了最大化聚合内存带宽，我们提供了具有并行协调的通信调度，以有效分组和路由具有相同关键数据路径的令牌。此外，我们提出了新的建模优化方法 PR-MoE 和 MoS，以在保持准确性的同时减小 MoE 模型大小。有关 DeepSpeed MoE 推理优化的更多信息，请参阅我们的博客文章。

DeepSpeed 为通过 DeepSpeed-MoE 库训练的 MoE 模型变体提供了无缝推理模式（MoE 教程）。为此，只需使用 DeepSpeed 推理引擎初始化模型，即可在评估模式下运行模型。

MoE 推理性能

在现代生产环境中，强大的深度学习模型通常使用数百个 GPU 设备进行服务，以满足流量需求并提供低延迟。探索如何在大规模 MoE 模型推理中实现高吞吐量和低延迟这两个宏伟目标至关重要。

对于密集模型，吞吐量可以通过使用多个 GPU 和数据并行（独立的副本，无 GPU 间通信）来提高，而通过张量切片等技术将模型分区到多个 GPU 来实现更低的延迟。在总吞吐量方面，最佳情况下的扩展是与 GPU 数量的增加呈线性关系，即每个 GPU 的吞吐量恒定。这对于纯数据并行推理场景是可能的，因为 GPU 之间没有通信。为了降低延迟，张量切片式的模型并行已被证明是有益的，但它也带来了成本——GPU 间的通信开销——这通常会降低每个 GPU 的吞吐量并导致总吞吐量的次线性扩展。换句话说，对于密集模型，我们无法同时利用并行性来优化延迟和吞吐量；它们之间存在权衡。然而，MoE 推理提供了独特的机会，可以在扩展到大量设备的同时，同时提供优化的延迟和吞吐量。

下图展示了我们如何同时实现低延迟和超线性吞吐量提升。我们在这篇论文中详细讨论了这一点。

52b-MoE-128

端到端 MoE 推理示例

在本节中，我们将通过一个端到端的示例，详细阐述 DeepSpeed 库中 MoE 推理支持的使用。

推理初始化

使用 DeepSpeed-MoE 进行推理时，请使用 init_inference API 来加载 DeepSpeed MoE 模型进行推理。在此，您可以指定模型并行/张量切片度（mp_size）、专家并行度（ep_size）以及专家数量（moe_experts）。我们根据 world_size（GPU 总数）和专家并行大小的最小值创建各种进程组。通过使用此组，我们可以在专家并行 GPU 之间划分专家。如果专家数量少于 GPU 总数，DeepSpeed-MoE 会利用专家切片在专家并行 GPU 之间划分专家参数。此外，如果模型尚未加载适当的检查点，您还可以使用 json 文件提供检查点描述，或者直接传递 'checkpoint' 路径来加载模型。要注入高性能推理内核，您可以将 replace_with_kernel_inject 设置为 True。


import deepspeed
import torch.distributed as dist

# Set expert-parallel size
world_size = dist.get_world_size()
expert_parallel_size = min(world_size, args.num_experts)

# create the MoE model
moe_model = get_model(model, ep_size=expert_parallel_size)
...

# Initialize the DeepSpeed-Inference engine
ds_engine = deepspeed.init_inference(moe_model,
                                     mp_size=tensor_slicing_size,
                                     dtype=torch.half,
                                     moe_experts=args.num_experts,
                                     checkpoint=args.checkpoint_path,
                                     replace_with_kernel_inject=True,)
model = ds_engine.module
output = model('Input String')

各种配置选项

在此，我们展示一个使用 MoE 模型进行文本生成的示例，其中我们可以指定模型并行大小和专家数量。DeepSpeed 推理引擎负责使用张量切片度、专家数量以及用于运行 MoE 模型的 GPU 总数来创建不同的并行组。关于专家参数，我们首先使用专家并行将每组专家分配给一个 GPU。如果 GPU 数量多于专家数量，我们则使用专家切片在 GPU 之间垂直/水平地划分每个专家。

我们来看看运行示例时传递的一些参数。请参考DeepSpeed 示例以获取完整的文本生成推理示例。

generate_samples_gpt.py \
       --tensor-model-parallel-size 1 \
       --num-experts ${experts} \
       --num-layers 24 \
       --hidden-size 2048 \
       --num-attention-heads 32 \
       --max-position-embeddings 1024 \
       --tokenizer-type GPT2BPETokenizer \
       --load $checkpoint_path \
       --fp16 \
       --ds-inference \

标准 MoE 模型性能

为了展示 DeepSpeed-MoE 推理随 GPU 数量增加而变化的性能扩展，我们考虑了一个具有 128 个专家和 1.3B 密集模型的 52B 模型架构，使用上述脚本中所示的参数。在此示例中，我们将张量切片度设置为 1，因为模型的非专家部分相对较小（805M 参数）。我们使用最后一个标志 ds-inference 来在 DeepSpeed-MoE 和 PyTorch 实现之间切换。

对于 DeepSpeed-MoE 推理，我们在本教程中展示了两个版本的结果：1) 通用版（Generic），即当前支持灵活并行和 PR-MoE 模型优化的 DeepSpeed 库开源版本；2) 专用版（Specialized），即 DeepSpeed MoE 推理系统最优化版本，包含将稍后发布的特殊计算和通信内核。正如我们在博客文章中提到的，MoE 推理优化将分阶段发布。

下图显示了在 8、16 和 32 个 GPU 上运行 PyTorch、DeepSpeed-MoE（通用版）和 DeepSpeed-MoE（专用版）三种不同配置的推理性能。与 PyTorch 相比，随着 GPU 数量的增加，DeepSpeed-MoE 获得了显著更高的性能优势。通过使用通用的 DeepSpeed-MoE 推理，我们可以获得比 PyTorch 高出 24% 到 60% 的性能提升。此外，通过启用 DeepSpeed-MoE 推理的全部功能，例如通信优化和 MoE 定制内核，性能加速得到进一步提升（2 倍 – 3.2 倍）。

52b-MoE-128

使用 PR-MoE 优化实现更快的性能和更低的推理成本

为了选择不同的 MoE 结构，我们在推理示例中添加了一个新参数，名为 mlp-type，用于在 'standard' MoE 结构和 'residual' 结构之间进行选择，以启用 PR-MoE 提供的建模优化。除了更改 mlp-type，在使用 PR-MoE 时，我们还需要以不同的方式传递专家数量。与每个 MoE 层使用相同数量专家的标准 MoE 不同，PR-MoE 对网络的初始层和更深层使用不同的专家数量。以下是 PR-MoE 的一个示例，它对每隔一层使用 64 个和 128 个专家的混合。

experts="64 64 64 64 64 64 64 64 64 64 128 128"
generate_samples_gpt.py \
       --tensor-model-parallel-size 1 \
       --num-experts ${experts} \
       --mlp_type 'residual' \
       --num-layers 24 \
       --hidden-size 2048 \
       --num-attention-heads 16 \
       --max-position-embeddings 1024 \
       --tokenizer-type GPT2BPETokenizer \
       --load $checkpoint_path \
       --fp16 \
       --ds-inference \

为了评估 PR-MoE 的性能，我们使用了两种模型结构，'standard' 和 'residual'，以及下表中所示的配置参数。由于我们无法将 24B+MoE-128 的非专家部分放入单个 GPU 中，因此我们使用了大于 1 的模型并行大小。我们选择张量切片度以获得最佳性能优势。

模型 大小（十亿） 层数 隐藏层大小 MP 度 EP 度
2.4B+MoE-128 107.7 16 3584 1 64 - 128
24B+MoE-128 1046.9 30 8192 8 64 - 128

我们使用 1 个节点（8 个 A100 GPU）对 2.4B+MoE-128 进行推理，并使用 8 个节点（64 个 A100 GPU）对 24B+MoE-128 进行推理。下图显示了三种不同配置的性能：MoE-Standard (PyTorch)、MoE-Standard (DeepSpeed-Generic) 和 PR-MoE (DeepSpeed-Generic)。通过使用标准 MoE，DeepSpeed 对这两个模型的推理性能分别比 PyTorch 提高了 1.4 倍和 1.65 倍。此外，通过使用 PR-MoE，我们可以将性能加速提高到 1.81 倍和 1.87 倍，同时保持模型质量不变。

52b-MoE-128

更多性能结果以及向更大模型和更多 GPU 扩展的数据可在我们的博客文章和论文中查看。

恭喜！您已完成 DeepSpeed MoE 推理教程。

更新日期：