DeepSpeed MoE 大规模稀疏模型推理实战指南：并行切分、PR-MoE 优化与 init_inference 全解析

在 DeepSpeed 中，基于 MoE（Mixture-of-Experts）训练的稀疏大模型可以通过 deepspeed.init_inference 无缝切换到推理模式，从而同时获得低延迟与超线性吞吐扩展。本文以仓库内教程 mixture-of-experts-inference.md 为骨架，结合 inference/engine.py 与 inference/config.py 源码，系统讲解 DeepSpeed MoE 推理所依赖的专家并行（expert-parallelism）、专家切片（expert-slicing）、张量切片（tensor-slicing）等机制，并给出可直接套用的 Python 与命令行推理示例，帮助读者完成从 MoE 模型加载、并行组配置到 PR-MoE 结构推理优化的完整实践。

MoE 推理的核心优化思路：多种并行度的组合

DeepSpeed-MoE Inference 是在稠密模型推理优化（DeepSpeed-Inference）基础之上，为大规模 MoE 模型引入的一组新特性。其基本思想是：把 MoE 模型中的不同参数按各自特性使用不同的并行策略切分，而不是对所有参数一刀切。

从 mixture-of-experts-inference.md 的介绍可以归纳出三条主线：

1. 非专家参数（attention、共享层等）使用数据并行（data-parallelism）与张量切片（tensor-slicing）；
2. 专家参数使用专家并行（expert-parallelism）与专家切片（expert-slicing）；
3. 通信调度与并行协调：为了最大化聚合内存带宽，DeepSpeed 会把共享同一条"关键数据路径"（critical-data-path）的 token 有效分组并路由，配合并行度做通信调度；
4. 建模级优化：提出 PR-MoE 与 MoS 两种模型优化手段，在不损失精度的前提下压缩 MoE 模型体积。

其中 PR-MoE 属于"可在当前开源版本直接使用"的模型结构优化，对应配置中的 mlp_type 'residual' 与逐层可变的专家数量（num_experts），我们会在后续章节给出完整命令示例。

对于使用 DeepSpeed-MoE 库训练的 MoE 模型变体，DeepSpeed 提供了"无缝推理"模式——无需改写模型结构，只需用 DeepSpeed-Inference 引擎对模型执行初始化并以 eval 模式运行即可（MoE 模型的训练教程见仓库中的 mixture-of-experts.md 与 mixture-of-experts-nlg.md）。

延迟与吞吐：为什么 MoE 推理能"鱼与熊掌兼得"

在讲解具体 API 之前，有必要先理解 MoE 推理的独特优势——它在扩展到大规模设备数时，可以同时优化延迟和吞吐，而这在稠密模型上通常是一对矛盾目标。

对稠密模型而言：

• 提升吞吐：靠多 GPU 数据并行（各 GPU 运行独立副本、无 GPU 间通信），吞吐理论上随 GPU 数线性增长，即每 GPU 吞吐恒定；
• 降低延迟：靠张量切片式模型并行（把模型切分到多 GPU），但会引入 GPU 间通信开销，降低单卡吞吐，导致总吞吐呈亚线性增长。

也就是说稠密模型难以同时用并行手段优化延迟与吞吐，二者存在权衡。而 MoE 推理提供了独特的优化空间，下图（出自教程正文）展示了如何在扩展设备数的同时实现低延迟与超线性吞吐增长（该结论的详细推导见教程引用的论文）：

图片出处：docs/assets/images/moe-lat-tput.png，正文引用原样保留于教程 mixture-of-experts-inference.md 的 "MoE Inference Performance" 一节。

用 init_inference 初始化 MoE 模型为推理模式

API 基本用法

对于推理，需要调用 DeepSpeed 顶层 API deepspeed.init_inference 加载 MoE 模型。初始化时需要显式指定三个与并行度相关的量：

• mp_size（tensor-slicing degree）：非专家部分的张量切片/模型并行度；
• ep_size（expert-parallelism degree）：专家并行度，即专家参数被切分到的 GPU 组大小；
• moe_experts（number of experts）：每个 MoE 层的全局专家总数。

教程给出的核心代码骨架如下（省略号部分为用户自己的模型装配逻辑）：

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

关键点解读：

• expert_parallel_size = min(world_size, args.num_experts)：这是基于世界大小（总 GPU 数）与专家并行度的最小值创建各种进程组的依据。在 world_size 大于专家数 num_experts 时，DeepSpeed 会退而使用专家切片（expert-slicing），把每个专家在专家并行的多个 GPU 间进行切分（垂直/水平切分权重）。
• checkpoint：如果模型尚未加载对应 checkpoint，既可以传入一个 json 描述文件，也可以直接传 checkpoint 路径。
• replace_with_kernel_inject=True：注入高性能推理 kernel。
• ds_engine.module：引擎包装后得到的、可直接前向推理的模型；model('Input String') 即为文本生成式调用。

init_inference 的 API 形态（源码确认）

从 deepspeed/init.py 的实现看，init_inference 的签名为：

def init_inference(model: torch.nn.Module,
                   config: Optional[Union[str, Dict[str, Any]]] = None,
                   **kwargs: Any) -> InferenceEngine:

它支持四种等价调用方式（源码注释中列出的 Case 1–4）：

1. 只传 model，使用默认配置；
2. 只传 config（字典或 JSON 文件路径）；
3. 只用关键字参数（如 tensor_parallel={"tp_size": world_size}, dtype=torch.half, replace_with_kernel_inject=True）；
4. config 与 kwargs 同时传入——两者会被合并，且 kwargs 优先级更高；若同名键取值不一致，会抛出 Conflicting argument ... in 'config' ... and kwargs ... 错误。

内部实现会先把 config（str 则读取 JSON）解析成字典，再与 kwargs 合并，最后构造 DeepSpeedInferenceConfig(**config_dict) 并创建 InferenceEngine（见 inference/engine.py）。

推理配置参数详解（含默认值与弃用映射）

为了兼容旧接口，教程示例中使用的顶层参数（mp_size、ep_size、moe_experts 等）在当前仓库的配置模型 DeepSpeedInferenceConfig 中已被标记为 deprecated，并映射到新的结构化配置字段。下表汇总了 MoE 推理相关参数及当前源码中的默认值（依据 inference/config.py）。

旧式顶层参数（教程写法）	新式结构化字段	默认值	说明
mp_size	tensor_parallel.tp_size	1	模型并行/张量切片度（见 config.py L278）
ep_size	moe.ep_size	1	专家并行度，用于把专家切分到专家并行组内的各 GPU（config.py L285）
moe_experts	moe.moe_experts（别名 num_experts）	[1]	MoE 层全局专家总数，支持逐层列表（config.py L298）
moe_type	moe.type	standard	MoE 结构类型，取值 standard / residual（config.py L299）
replace_with_kernel_inject	同名（别名 kernel_inject）	False	置真注入高性能推理 kernel
dtype	同名	torch.float16	模型数据类型，支持 torch.half/torch.float/torch.int8/bf16 等
checkpoint	同名	None	DeepSpeed 兼容 checkpoint 路径，或带加载策略的 JSON
—	triangular_masking	True	控制 attention 得分掩码类型

新式结构化配置实际定义于 DeepSpeedMoEConfig，包含：

• enabled: bool = True：是否启用 MoE；
• ep_size: int = 1：专家并行度；
• moe_experts: list = [1]（别名 num_experts）：MoE 层使用的全局专家数，用列表表示即可支持 PR-MoE 这种逐层不同专家数的结构；
• type: MoETypeEnum = standard：MoE 层类型，枚举值 standard / residual（对应 MoETypeEnum）；
• ep_mp_group 与 ep_group（别名 expert_group）：预置的进程组对象。

张量并行配置见 DeepSpeedTPConfig：tp_size（默认 1）、tp_grain_size（默认 64，DNN 库更偏好 2 的幂粒度）、mpu（模型并行单元对象，需实现 get_{model,data}_parallel_{rank,group,world_size}()）、tp_group。

引擎如何创建专家并行组（源码行为）

InferenceEngine 初始化时会先通过 has_moe_layers 探测模型是否含 MoE 层（见 deepspeed/moe/utils.py），当 moe and dist.get_world_size() > 1 时调用 _create_ep_parallel_group(config.moe.ep_size)（engine.py L132-L133），其逻辑（engine.py L291-L318）与教程描述一一对应：

• 按每个 ep_size 建立 ep_group：num_ep_groups = world_size // ep_size，当 ep_size > world_size 时组内实际使用 world_size 个 rank；
• 当 world_size > ep_size 时，额外构建 expert_mp_group（专家模型并行组），用于专家切片——每个专家在多个 GPU 间进一步切分，这正是教程所说"如果 GPU 数大于专家数，就用 expert-slicing 把每个专家垂直/水平切到多张卡上"的底层实现。

换句话说，DeepSpeed 推理引擎会根据张量切片度、专家数、总 GPU 数三者的关系，自动完成不同并行度进程组的创建，无需用户手动管理通信组。

端到端文本生成示例：命令行参数详解

DeepSpeed 推理引擎负责按张量切片度、专家数与总 GPU 数创建不同的并行组；对专家参数，先用专家并行把每组专家指派到一块 GPU，若 GPU 数多于专家数，再用专家切片把每个专家跨 GPU 垂直/水平切分。仓库内对应的完整 generate-text 推理示例脚本（Megatron-DeepSpeed 集成示例 generate_text.sh，详见教程中引用的 DeepSpeed-Example）基于 generate_samples_gpt.py，典型调用如下：

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 \

参数含义速查：

参数	作用
--tensor-model-parallel-size	张量/模型并行度，即上文 mp_size
--num-experts	每层专家数；支持以空格分隔的多个值用于 PR-MoE 逐层结构
--num-layers / --hidden-size / --num-attention-heads	Transformer 结构规模：层数、隐藏维、注意力头数
--max-position-embeddings	最大序列长度相关的位置编码上限
--tokenizer-type	分词器类型（此处为 GPT-2 BPE）
--load $checkpoint_path	训练好的 checkpoint 路径
--fp16	半精度运行（与 dtype=torch.half 对应）
--ds-inference	关键开关：切换 DeepSpeed-MoE 与 PyTorch 原生实现

示例中 24 层、2048 隐藏维、32 注意力头对应一个约 1.3B 的稠密 backbone，其非专家部分相对较小。

用 PR-MoE 结构进一步提速并降低推理成本

在标准 MoE 与 PR-MoE 两种结构之间切换，需要：

1. 新增参数 --mlp-type：取值 'standard'（标准 MoE）或 'residual'（PR-MoE）；
2. 以不同方式传入专家数：标准 MoE 每一层使用相同专家数；PR-MoE 中网络较浅层与较深层使用不同的专家数量。

下面是一个"每两层交替使用 64 与 128 个专家"的 PR-MoE 示例（24 层网络，共 12 组专家数，即隔层 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 \

在配置模型层面，"每层专家数不同"正是通过 num_experts / moe_experts 的列表语义表达；而 'residual' 结构对应 MoETypeEnum 中的 residual 取值。

PR-MoE 评估所用的模型规模

为了评估 PR-MoE，教程采用 'standard' 与 'residual' 两种结构，配置如下表。由于 24B+MoE-128 的非专家部分无法放入单卡，其张量并行度大于 1；选择张量切片度以获得最佳性能收益：

Model	Size (billions)	#Layers	Hidden size	MP degree	EP degree
2.4B+MoE-128	107.7	16	3584	1	64 - 128
24B+MoE-128	1046.9	30	8192	8	64 - 128

注：教程同时提到"模型质量保持"这一前提——PR-MoE 在不牺牲精度的前提下把加速比进一步提升。

多 GPU 扩展下的推理性能报告

标准 MoE 模型（52B MoE-128）的扩展性

为展示 DeepSpeed-MoE 推理随 GPU 数增长的可扩展性，教程选用了 128 个专家、约 1.3B 稠密骨干的 52B MoE 结构（如上节脚本参数）。由于非专家部分相对较小（约 805M 参数），示例将张量切片度设为 1，仅用专家并行与专家切片扩展。

针对 DeepSpeed-MoE 推理，教程区分两个版本：

• Generic（通用版）：当前开源版本，支持灵活并行度与 PR-MoE 模型优化；
• Specialized（专用版）：DeepSpeed MoE 推理系统的最优化版本，含专门的计算与通信 kernel（教程发布时计划后续分批开源）。

下图为 PyTorch、DeepSpeed-MoE（Generic）、DeepSpeed-MoE（Specialized）三种配置在 8/16/32 GPU 上的推理性能对比：

据教程记录：相比 PyTorch，Generic 版 DeepSpeed-MoE 推理可带来约 24%–60% 的性能提升；开启通信优化与 MoE 定制 kernel（Specialized 版）后，加速比可进一步提升至 2x–3.2x。图片文件位于 docs/assets/images/1.3B-MoE-128.png。

PR-MoE 的进一步加速效果

教程用 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.4x 与 1.65x；改用 PR-MoE 后，加速比进一步提升到约 1.81x 与 1.87x，同时保持模型质量。图片文件位于 docs/assets/images/prmoe.png。

说明：以上性能数据均来自该教程正文及其引用论文中的特定硬件环境（A100 GPU、特定模型规模），属于官方文档报告的实验结论；实际加速效果会随模型结构、GPU 型号、专家数与并行度设置而变化，应以本机实测为准。

从源码理解推理链路：一份代码导读

为了让读者把"教程用法"落到"代码实现"上，这里给出与 MoE 推理直接相关的几个关键源码位置：

1. 顶层入口 deepspeed/init.py：init_inference 的四种调用形态、config 与 kwargs 合并及冲突校验，最终返回 InferenceEngine。
2. 推理引擎 deepspeed/inference/engine.py：初始化流程包括 dtype 转换、TP 组创建（_create_model_parallel_group，L280-L289）、MoE 层探测与 EP/专家切片进程组创建（L127-L133、_create_ep_parallel_group L291-L318）、注入策略执行（用户注入策略 / kernel 注入 / 自动 TP 三种模式，L135-L169）。
3. 配置模型 deepspeed/inference/config.py：DeepSpeedMoEConfig、DeepSpeedTPConfig、DeepSpeedInferenceConfig 及全部弃用别名映射，是解释旧教程参数与新配置字段对应关系的权威依据。
4. checkpoint 加载：inference/engine.py 的 _load_checkpoint/_get_all_ckpt_names 支持读取以 mp_rank_XX_model_states.pt 形式保存的模型分片（读取 latest 标签定位最新版本），也支持通过 JSON 策略文件描述 checkpoint，这与 init_inference 中 checkpoint 参数的两种传法一致。
5. MoE 结构：deepspeed/moe/utils.py 的 has_moe_layers 负责在引擎初始化时判定模型是否包含 MoE 层，决定是否触发上述专家并行/切片组的创建。

小结

通过本文可以掌握三条关键链路：

1. 并行切分：ep_size 决定专家并行组，专家数小于 GPU 数时自动启用专家切片；mp_size/tensor_parallel.tp_size 负责非专家部分的多卡切分；
2. 配置映射：旧的顶层参数（mp_size、ep_size、moe_experts）在当前仓库中已映射为 tensor_parallel.* 与 moe.* 结构化配置，迁移代码时注意兼容与弃用提示；
3. 建模优化：在标准 MoE 之外，可通过 --mlp_type 'residual' 与逐层专家数列表启用 PR-MoE，在保持模型质量的同时获得更高推理吞吐。

完成以上步骤，即可将 DeepSpeed-MoE 训练的模型平滑切换为推理模式并充分享受多 GPU 扩展带来的吞吐收益。