使用 DeepSpeed-MoE 进行大规模 MoE 模型推理入门

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

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

MoE 推理性能Permalink

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

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

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

52b-MoE-128

端到端 MoE 推理示例Permalink

在本部分中，我们将详细说明使用端到端示例在 DeepSpeed 库中使用 MoE 推理支持。

初始化推理Permalink

对于 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')

各种配置选项Permalink

在这里，我们展示了一个使用 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 模型的性能Permalink

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

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

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

52b-MoE-128

使用 PR-MoE 优化获得更快的性能和更低的推理成本Permalink

为了在不同的 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 将两个模型的推理性能分别提高了 1.4 倍和 1.65 倍，与 PyTorch 相比。此外，通过使用 PR-MoE，我们可以将性能加速提高到 1.81 倍和 1.87 倍，同时保持模型质量。

52b-MoE-128

更多性能结果以及扩展到更大模型和更多 GPU 的内容，请参阅我们的博客文章和论文。

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

更新：