DeepSpeed Inference 内核优化深度解析：多 GPU 自适应并行、算子融合推理内核与 MoQ 量化支持

本篇技术解读围绕 DeepSpeed 官方博客《Multi-GPU inference with customized inference kernels and quantization support》展开。该博客系统阐述了 DeepSpeed Inference 如何针对大模型推理场景中的三大瓶颈——多 GPU 部署、小批量推理的 kernel 性能与量化支持——给出三条技术路线：推理自适应并行（inference-adapted parallelism）、面向 Transformer 块的定制化推理内核（customized inference kernels） 与 MoQ 量化工具包。读者在读完本文后，将理解 DeepSpeed Inference 从训练检查点到多 GPU 低延迟推理的完整技术链路，掌握算子融合推理内核的底层设计思路、自动 kernel 注入的使用方式，以及 MoQ 量化的训练与推理配合方案。

本文以该博客为主体骨架，并对照当前仓库中 deepspeed/ 目录下的实际实现、docs/_tutorials/inference-tutorial.md 教程与 docs/_posts/2021-05-05-MoQ.md 进行交叉印证，帮助你在阅读历史设计的同时，快速定位仓库内对应的源码与可运行示例。

一、背景：大模型"训得好"不代表"用得顺"

DeepSpeed 在训练侧已经支持超大规模模型的分布式训练，但把训练好的模型真正部署到目标应用场景仍然面临三大限制（这也是 DeepSpeed Inference 立项要解决的核心问题）：

1. 缺少多 GPU 推理支持：大模型单卡放不下，且存在延迟（latency）要求，需要跨设备协同推理能力；
2. 小批量推理下 kernel 性能受限：推理请求往往是低 batch 的在线负载，通用 GPU kernel（如 cuBLAS GeMM）在小 batch 场景下利用率不高；
3. 量化难以利用：既缺少把模型量化以压缩体积与降低延迟的完整流程，也缺少无需专用硬件就能高性能运行量化模型的推理 kernel。

针对上述挑战，DeepSpeed Inference 提出了三大关键能力：推理自适应并行（多 GPU 推理）、面向小 batch 调优的推理专用内核，以及同时覆盖量化感知训练与量化模型推理内核的灵活量化支持。这三者相互配合，构成了从训练权重到生产推理的完整通路。

二、推理自适应并行：为延迟与部署成本重塑 MP/PP

并行是把大模型装进显存、降低单设备内存消耗的有效手段，训练与推理皆是如此。但把训练时的并行配置直接搬到推理并不合适：

• 训练阶段的模型并行（MP）与流水并行（PP）度数，通常是依据训练时的显存占用、计算风格与资源预算在启动前就固定下来的；
• 而推理的计算本质对显存需求更低，因此可以承担更大的单设备分片，这意味着部署时可以减少所需的并行度数；
• 训练优化的核心指标是吞吐，推理则把**延迟（latency）**当作一等公民。

因此 DeepSpeed Inference 的做法是：自动把 MP 作为降低模型延迟的主要手段，并优先确定其并行度数。通过 MP 把模型切分并跨多 GPU 并行计算，可以降低延迟；但同时它也会缩小计算粒度、增加通信，可能损害吞吐。当延迟目标满足之后，DeepSpeed 再施加流水并行（PP）来最大化吞吐。整体上，DeepSpeed Inference 支持从训练到推理时对并行方式与并行度数二者的灵活适配，在降低延迟的同时节省部署成本。

这一能力在教程中有明确的落地形态。当前仓库的 inference-tutorial.md 指出："DeepSpeed supports running different MP degree for inference than from training. For example, a model trained without any MP can be run with MP=2, or a model trained with MP=4 can be inferenced without any MP. DeepSpeed automatically merges or splits checkpoints during initialization as necessary." 即推理时可以使用与训练不同的 MP 度数，初始化过程中 DeepSpeed 会自动完成 checkpoint 的合并或切分。

2.1 通过 init_inference 指定推理并行配置

推理时只需通过 deepspeed.init_inference 声明模型并行规模即可，入口位于 deepspeed/init.py。以 GPT-Neo-2.7B 多卡生成为例（完整脚本见 inference-tutorial.md）：

# gpt-neo-2.7b-generation.py
import os
import deepspeed
import torch
from transformers import pipeline

local_rank = int(os.getenv('LOCAL_RANK', '0'))
world_size = int(os.getenv('WORLD_SIZE', '1'))
generator = pipeline('text-generation', model='EleutherAI/gpt-neo-2.7B',
                     device=local_rank)

generator.model = deepspeed.init_inference(generator.model,
                                           tensor_parallel={"tp_size": world_size},
                                           dtype=torch.float,
                                           replace_with_kernel_inject=True)

string = generator("DeepSpeed is", do_sample=True, min_length=50)
if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0:
    print(string)

使用 DeepSpeed 启动器在两张 GPU 上运行：

deepspeed --num_gpus 2 gpt-neo-2.7b-generation.py

值得说明的是，这里的模型本身是单卡 checkpoint、未经过任何模型并行训练，却可以在推理时直接做跨卡 tensor-slicing，这正是"推理自适应并行"能力的直接体现。

2.2 跨训练框架的 checkpoint 加载约定

从当前仓库的 inference-tutorial.md 可以看到，Megatron-LM 以 MP=2 训练的多卡 checkpoint 需要以 JSON 描述清单：

"checkpoint.json":
{
    "type": "Megatron",
    "version": 0.0,
    "checkpoints": [
        "mp_rank_00/model_optim_rng.pt",
        "mp_rank_01/model_optim_rng.pt",
    ],
}

而对于 DeepSpeed 训练产出的 checkpoint，JSON 只需给出 checkpoint 路径：

"checkpoint.json":
{
    "type": "ds_model",
    "version": 0.0,
    "checkpoints": "path_to_checkpoints",
}

在 init_inference 中通过 checkpoint 参数传入该 JSON（或在 pre_load_checkpoint 场景下直接传入已加载好的模型），DeepSpeed 会自动完成模型切分、kernel 注入与卡间通信管理。

三、面向 Transformer 块的定制化推理内核：四大融合区域

为了在高计算效率上取得收益，DeepSpeed Inference 提供了针对 Transformer 块的推理内核，且在设计之初就考虑了多 GPU 场景下的模型并行。它与同类 kernel-fusion 方案的核心差异在于：不仅融合逐元素操作（element-wise operations，如 bias-add、残差连接、激活函数），还把通用矩阵乘（GeMM）与其他操作合并在一起。为此，DeepSpeed 设计了高效的向量-矩阵（vector-matrix）或瘦矩阵-矩阵（skinny matrix-matrix）乘法实现，使得可以在 GeMM 的 reduction 边界融合更多的算子。

3.1 融合的两大设计原则

算子融合遵循两条主要策略：

1. 在整条被融合的算子链上保持输入/输出的访问模式（access-pattern）不变。这样可以保证不同 thread-block 之间不会发生跨 Streaming-Multiprocessor（SM）的数据搬运——SM 之间除了经由主存的间接通信外没有直接的通信手段，而主存通信会因内存访问的非确定性行为引入 block 级同步开销。
2. 在每一个 all-reduce 边界上进行融合。因为在模型并行下，部分结果必须在各 GPU 之间完成归约之后才能继续执行，all-reduce 天然是融合操作的"同步点"。

3.2 一个 Transformer 层中的四大融合区域

博客中的 Figure 1 给出了 Transformer 层各组件与推理优化中考虑的融合分组（虚线宽度表示融合深度），并明确采用 NVIDIA Megatron-LM 风格的并行方式——把注意力（Attn）与前馈（FF）块切分到多张 GPU，从而在 Attn 与 FF 块之后各包含一次跨 GPU 的 all-reduce。

围绕一个 Transformer 层，推理内核在四个主要区域进行算子融合：

1. 输入 Layer-Norm + Q、K、V GeMM 及其 bias-add：把归一化、多头投影与偏置加法串成一条融合链；
2. Transform 与 Attention 的合并：使用**隐式矩阵变换（implicit matrix transformation）**合并注意力计算涉及的多个 GeMM，以降低中间内存压力；
3. 中间 FF + Layer-Norm + Bias-add + 残差 + GELU：这是融合深度最大的区域，依赖新的 GeMM 调度把尽可能多的操作接续起来；
4. Bias-add + 残差：输出侧的轻量融合。

3.3 底层的实现机制

在具体实现上，这些融合依赖于三样关键机制：

• **共享内存（shared memory）**作为中间缓存，在 Layer-Norm/GeMM 所用的归约操作与逐元素操作之间传递数据；
• **warp 级指令（warp-level instructions）**用于归约部分计算结果时在线程间通信；
• 为 GeMM 设计的新调度（schedule），它允许在第三个 kernel-fusion 中按需接入尽量多的操作。

博客同时给出了对照数据：与使用 cuBLAS GeMM 的非融合计算方式相比，上述四类 kernel-fusion 分别带来约 1.5x、2.9x、3x 与 1.2x 的性能提升。

3.4 在源码中的落地位置

如果你希望在仓库中追踪这套推理内核的实际实现，可以沿以下路径展开：

• CUDA/C++ 推理算子源码位于 csrc/transformer/inference，覆盖 GEMM、归一化、Softmax、激活与量化相关 kernel 实现；
• 对应的 Python 封装与 builder 见 op_builder/transformer_inference.py 与 op_builder/inference_core_ops.py；
• 推理算子的 Python 侧接入逻辑分布在 deepspeed/ops/transformer/inference；
• 推理时算子的装载与图替换逻辑位于 deepspeed/inference/engine.py。

四、从训练到推理的无缝流水线：自动内核注入与 Policy 映射

运行推理模式时，DeepSpeed Inference 只要求两样东西：模型 checkpoint 的位置与期望的并行配置（MP/PP 度数），其余工作（模型自动切分、kernel 注入、卡间通信管理）都由框架完成。

对许多已知模型架构，推理内核可以通过预定义的 policy map 直接启用——policy 会建立"用户原始层参数 ↔ 推理内核参数"的映射关系。博客列举了 HuggingFace（BERT、GPT-2）与 Megatron GPT 系列模型；对其他 Transformer 模型，用户可以自行编写 policy map。要点有二：

1. 只要拿到全部模型 checkpoint，DeepSpeed Inference 可以独立于训练管线运行；
2. 只要定义了正确的映射 policy，DeepSpeed Transformer 推理内核就能注入到任意 Transformer 模型中。

4.1 内置支持的模型族（源码印证）

在 deepspeed/module_inject/replace_policy.py 中可以看到当前仓库内置支持的注入策略列表：

# transformer-based policies
replace_policies = [
    HFBertLayerPolicy, HFGPTNEOLayerPolicy, GPTNEOXLayerPolicy, HFGPTJLayerPolicy, MegatronLayerPolicy,
    HFGPT2LayerPolicy, BLOOMLayerPolicy, HFOPTLayerPolicy, HFCLIPLayerPolicy, HFDistilBertLayerPolicy,
    LLAMALayerPolicy, LLAMA2LayerPolicy, InternLMLayerPolicy
]

各模型的具体参数映射策略实现位于 deepspeed/module_inject/containers 目录。相比 2021 年博客发布时的 BERT/GPT-2/Megatron 时代，如今的内置覆盖已扩展到 BLOOM、GPT-Neo、GPT-J、GPT-NeoX、OPT、CLIP、Llama 与 InternLM 等架构。

4.2 内核自动注入的开启方式

在 deepspeed/inference/config.py 的 DeepSpeedInferenceConfig 中，可以找到与自动注入直接相关的字段：

• replace_with_kernel_inject（别名 kernel_inject）：置为 True 即对兼容模型注入高性能推理内核；
• injection_policy（别名 injection_dict）：对框架尚不支持内核、或希望自行指定替换关系的模型，传入 {模块类型: (该层参与跨卡 all-reduce 的线性层参数路径)}；
• dtype：默认 torch.float16，可选 fp32 / fp16 / int8；
• tensor_parallel（别名 tp）：以 {"tp_size": world_size} 形式声明模型并行度数；
• checkpoint / checkpoint_dir：推理 checkpoint 的加载策略与目录；
• max_out_tokens / min_out_tokens（别名 max_tokens / min_tokens）：推理输出长度的上下界（默认 1024 / 1）。

以下是教程中的典型用法（HuggingFace pipeline 场景下仅替换模型即完成内核注入）：

# Initialize the DeepSpeed-Inference engine
ds_engine = deepspeed.init_inference(model,
                                     tensor_parallel={"tp_size": world_size},
                                     dtype=torch.half,
                                     checkpoint=None if args.pre_load_checkpoint else args.checkpoint_json,
                                     replace_with_kernel_inject=True)
model = ds_engine.module

对于不提供内核的模型，只使用模型并行时，需要显式给出 injection_policy，指明 Transformer Encoder/Decoder 层上两个参与跨卡通信的线性层——1) 注意力输出 GeMM、2) 层输出 GeMM。这是因为只有在这两处插入 all-reduce，才能把各模型并行 rank 上的部分结果正确合并。教程中以 T5 为例：

pipe = pipeline("text2text-generation", model="google/t5-v1_1-small", device=local_rank)
pipe.model = deepspeed.init_inference(
    pipe.model,
    tensor_parallel={"tp_size": world_size},
    dtype=torch.float,
    injection_policy={T5Block: ('SelfAttention.o', 'EncDecAttention.o', 'DenseReluDense.wo')}
)

说明：博客中提到的"推理内核注入 + 并行度指定"完整操作步骤，在仓库中有配套可运行的端到端示例，即 inference-tutorial.md，其中包含 HuggingFace pipeline 与 DeepSpeed 推理结合的 GPT-Neo-2.7B 文本生成完整客户端代码与启动命令。

五、DeepSpeed Quantization Toolkit 与 MoQ：无痛量化 + 高效量化推理

为进一步削减大模型推理成本，博客引入了 DeepSpeed Quantization Toolkit（量化工具包），它同时覆盖两个层面：灵活的量化感知训练（quantize-aware training） 与 面向量化模型的高性能推理内核。

5.1 训练侧：Mixture of Quantization（MoQ）

MoQ 是一套受混合精度训练启发、把量化"无缝"施加到训练流程中的新方法：

• 在训练每一步更新参数时，通过模拟量化带来的影响来控制模型精度；
• 支持灵活的量化策略与调度（quantization policies and schedules）——博客指出，在训练过程中动态调整量化比特数，最终得到的量化模型在相同压缩率下精度更高；
• 为适配不同任务，MoQ 还能利用模型的二阶信息检测其对精度的敏感性，据此调整量化调度与量化目标。

关于 MoQ 的独立深度解读可参阅仓库内同期博客 docs/_posts/2021-05-05-MoQ.md，其配套实战教程见 docs/_tutorials/MoQ-tutorial.md。

5.2 推理侧：面向量化模型的高性能内核

为了把量化模型的性能收益放到最大，DeepSpeed Inference 提供专门为量化模型设计的推理内核——它通过优化数据搬运来降低延迟，且不依赖专用硬件（无需为 INT8 推理配备定制加速芯片）。同时，工具包在客户端侧不要求任何代码改动，使用门槛很低。

量化推理在使用层面的体现是：推理时选择 dtype=torch.int8，并把 MoQ 量化时使用的分组设置（量化组数 quantize_groups、MLP 部分是否做额外分组 mlp_extra_grouping）通过 quantization_setting 传入 init_inference：

import deepspeed
model = deepspeed.init_inference(model,
                                 checkpoint='./checkpoint.json',
                                 dtype=torch.int8,
                                 quantization_setting=(quantize_groups,
                                                       mlp_extra_grouping)
                                )

对应到当前仓库的 deepspeed/inference/config.py，dtype 字段即承载 fp32 / fp16 / int8 三种取值，DeepSpeed 会依据该字段选择为该数据类型调优过的内核。量化算子的 CUDA 实现可在 csrc/quantization 与 csrc/transformer/inference 中找到线索。

六、性能结果：吞吐、GPU 资源与延迟的三重收益

6.1 吞吐提升与推理成本下降

博客报告了对应三个量级 Transformer 网络——GPT-2、Turing-NLG 与 GPT-3（同规模特征模型）的单 GPU 推理吞吐。与基线相比：

• 在与基线相同的 FP16 精度下，DeepSpeed Inference 将单 GPU 吞吐提升了 2～4 倍；
• 在 FP16 基础上叠加量化后吞吐进一步提升：GPT-2 提升约 3x、Turing-NLG 提升约 5x、与 GPT-3 特征和规模相近的模型提升约 3x；
• 上述吞吐提升直接换算为托管这些大模型的推理成本下降 3～5 倍；
• 并且这些吞吐与成本收益不以牺牲延迟为代价（见下文的 Figure 5）。

6.2 托管 GPU 数量的大幅缩减

推理成本下降的另一来源是托管大模型所需 GPU 数量减少，其优化来源有二：

1. 推理自适应并行：允许用户依据训练好的 checkpoint 调整模型并行与流水并行度数；
2. INT8 量化（MoQ）：把模型内存占用压缩一半。

博客报告（Figure 4）：通过自适应并行，17B 模型推理所用 GPU 减少一半；再叠加 MoQ INT8 量化后，17B 与 175B 模型分别减少到 1/4 与 1/2 的 GPU 用量。

6.3 延迟优化

对于延迟敏感的在线场景，可以进一步提高推理时的模型并行度数来压低延迟。博客报告（Figure 5）：

• 在模型并行度数提高到 4 时，与 PyTorch 基线相比延迟降低约 2.3x；
• 通过推理时的并行适配 + MoQ 量化，即使在使用远少于基线的 GPU 时依然能获得明显延迟改善：分别在使用 1/4 与 1/2 资源的条件下取得约 1.3x 与 1.9x 的加速。

以上性能数据均出自 DeepSpeed 官方博客在 2021 年发布时的基准测试报告，属于该项目自身的阶段性测量结果；具体数值会随硬件、模型与版本环境变化，引用时请注意其适用前提。

七、结语：从 2021 的设计看今天的 DeepSpeed Inference

回到核心结论，DeepSpeed Inference 的三大支柱——推理自适应并行、面向 Transformer 块的融合推理内核与 MoQ 量化的训练/推理全流程支持——在最初发布时解决的是三件事：让放不下的大模型跑起来、让跑起来的延迟可控、让跑起来的成本更低。

对照当前仓库可以发现，这套设计思想至今仍是 DeepSpeed 推理体系的基石：

• init_inference 入口与其配置对象 DeepSpeedInferenceConfig 保留在 deepspeed/init.py 与 deepspeed/inference/config.py 中，参数形态已演化为 tensor_parallel={"tp_size": ...} 的新式写法；
• 内核注入与策略映射体系沉淀于 deepspeed/module_inject（含 replace_policy.py 与 containers/）；
• 仓库中另存在新一代推理实现子包 deepspeed/inference/v2，代表后续持续演进的推理引擎（仓库亦发布了 DeepSpeed-FastGen 系列博客），但"从训练 checkpoint 无缝衔接、多 GPU 自适应并行、专用内核注入与量化支持"这一原始设计主线始终贯穿其中。

如果你希望深入动手实践，仓库中的 inference-tutorial.md 提供了从 init_inference 到 HuggingFace pipeline 的端到端脚本，MoQ-tutorial.md 则覆盖量化训练与 int8 推理的配置细节，二者是与本文配套的最佳下一步阅读材料。