DeepSpeed Inference 实战指南：用 init_inference 与内核注入高效服务 Transformer 模型

DeepSpeed-Inference 已全面升级为 DeepSpeed-FastGen，如需最新的性能优化、功能特性与模型支持，可优先阅读仓库内的 DeepSpeed-FastGen 发布博客。本文聚焦经典 init_inference API，讲解其模型并行、内核注入与量化推理的完整用法。

DeepSpeed-Inference 是 DeepSpeed 面向 Transformer 类 PyTorch 模型推理场景提供的一套无缝加速方案，解决的核心痛点是：模型太大装不进单卡、延迟与成本太高、以及量化后精度/性能难两全。它不需要你改动任何建模代码或导出专门的 checkpoint，即可让基于 DeepSpeed、Megatron 或 HuggingFace 训练的兼容模型以"零改造"方式跑起多卡推理。读完本文，你将掌握如何通过 deepspeed.init_inference 初始化推理引擎、使用模型并行（Tensor Parallelism）加载超大模型、自动注入高性能推理内核、按 Megatron/DeepSpeed 训练格式加载 checkpoint，以及如何用 fp32/fp16/int8（含 MoQ 量化）不同数据类型完成端到端的文本生成服务。

DeepSpeed-Inference 的核心能力与设计理念

DeepSpeed-Inference 为高效服务 Transformer 模型引入了三类关键能力：

1. 模型并行（Model Parallelism，MP）：当模型规模超出单卡 GPU 显存时，可以把模型切分到多张 GPU 上；即便对小模型，MP 也能通过降低单卡计算量来进一步压低推理延迟。
2. 推理专用内核（Inference-customized Kernels）：针对注意力、MLP 等算子进行深度优化，进一步减少延迟与成本。
3. 混合量化方法 MoQ（Mixture-of-Quantization）：在生产环境中既压缩模型体积、又降低推理开销。

在兼容模型的支持方式上，DeepSpeed 提供无缝推理模式：只要模型是由 DeepSpeed、Megatron 或 HuggingFace 训练得到的 transformer 模型，就无需导出模型、也无需从训练 checkpoint 转换出另一种格式。你需要做的只是在多卡推理时给出**模型并行度（MP degree）**与 checkpoint 信息（或一个已经加载好权重的模型），剩下的工作——按需切分模型、把兼容的高性能内核注入模型、管理 GPU 间通信——全部由 DeepSpeed 自动完成。

初始化推理引擎：init_inference 的四种调用形态

推理的统一入口是 init_inference，其定义位于 deepspeed/init.py：

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

model 是原始的、未加任何包装的 nn.Module 对象；config 可选，可以是一个 DeepSpeed 推理配置字典，也可以是某个 JSON 配置文件的路径。从源码实现看，配置与关键字参数（kwargs）会先合并（kwargs 覆盖同名配置），若同一键在两边取值冲突则会抛出 Conflicting argument 错误，随后统一交给 DeepSpeedInferenceConfig 完成校验，并构造 InferenceEngine。因此官方支持四种等价用法：

• Case 1：不传任何配置，全部使用默认配置：deepspeed.init_inference(generator.model)；
• Case 2：只传 config（字典或 JSON 路径）；
• Case 3：只传 kwargs，例如 tensor_parallel、dtype、replace_with_kernel_inject；
• Case 4：config 与 kwargs 同时给出，两者合并、kwargs 优先。

调用 init_inference 时你可以指定 MP 大小；如果模型尚未用对应 checkpoint 加载权重，还可以通过 JSON 描述文件或 checkpoint 路径来提供权重来源。返回的 InferenceEngine 包装了模型，其 .module 属性即替换/切分后的模型，可直接放入 HuggingFace pipeline 使用。

下面是一段典型的初始化代码骨架：

# create the model
if args.pre_load_checkpoint:
    model = model_class.from_pretrained(args.model_name_or_path)
else:
    model = model_class()

# create the tokenizer
tokenizer = tokenizer_class.from_pretrained(args.model_name_or_path)
...

import deepspeed

# 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
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
output = pipe('Input String')

需要特别留意的是：checkpoint=None if args.pre_load_checkpoint else args.checkpoint_json 表达的是"若已用 from_pretrained 预加载权重则不再提供 checkpoint，否则传入 checkpoint 描述"这一逻辑，是演示性写法，请在实际脚本中按自己的加载方式替换。

自动内核注入：replace_with_kernel_inject

要对兼容模型注入高性能内核，需要把 replace_with_kernel_inject 设为 True。从 inference/config.py 可以看到该参数支持别名 kernel_inject，默认值为 False；官方在参数注释中给出的典型适配模型包括 Bert、GPT-2、GPT-Neo、GPT-J 等。

**哪些模型"兼容"？**兼容策略由 deepspeed/module_inject/replace_policy.py 统一登记。从该文件源码看，replace_policies 列表覆盖了 HFBert、HFGPTNeo、GPTNeoX、HFGPTJ、Megatron、HFGPT2、BLOOM、HFOPT、HFCLIP、HFDistilBert、LLaMA、LLaMA2 与 InternLM 等 transformer 类策略，另有 UNet、VAE 两类非 transformer 策略登记在 generic_policies 中。如果你的模型不在支持列表内，可以向社区提交 PR，在 replace_policy 相关策略类中定义一个新的 policy，说明 Transformer 层的注意力与前馈部分各自参数如何映射。

Policy 类的本质，是把用户自定义 layer 实现的参数，与 DeepSpeed 推理优化过的 Transformer 层之间建立一一映射关系。策略类对应的容器实现位于 deepspeed/module_inject/containers（如 HFGPT2LayerPolicy、HFBertLayerPolicy、BLOOMLayerPolicy 等）。

手动注入策略：对不支持内核的模型只做模型并行

如果我们不想（或还不能）对某个模型注入内核，只想借助模型并行把模型摊到多张 GPU 上，可以传入一个 injection_policy，指明 Transformer Encoder/Decoder 层中两个关键的线性层：

1. 注意力输出 GeMM（attention output projection）；
2. 层输出 GeMM（transformer layer output projection）。

这两个位置是模型并行时必须插入 GPU 间 all-reduce 通信、以合并各 MP rank 上部分结果的地方。下面的例子演示了对 T5 模型使用 DeepSpeed-Inference 的方式：

# create the model
import transformers
from transformers.models.t5.modeling_t5 import T5Block

import deepspeed

pipe = pipeline("text2text-generation", model="google/t5-v1_1-small", device=local_rank)
# Initialize the DeepSpeed-Inference engine
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')}
)
output = pipe('Input String')

这里把 HuggingFace 的 T5Block 映射到三个参数名上：自注意力的 SelfAttention.o、跨注意力的 EncDecAttention.o，以及位置无关的 DenseReluDense.wo。在 inference/config.py 中，该参数对应的配置字段为 injection_policy（别名 injection_dict），格式为 {客户 nn.Module: 其对应的注入策略} 的字典。

在推理引擎侧，deepspeed/inference/engine.py 明确把工作模式归结为三类：① 用户显式指定的 injection policy（纯 tensor parallelism）；② 自动内核注入（replace_with_kernel_inject=True）；③ tp_size > 1 时的自动 tensor parallelism。同时引擎会断言 injection_policy 与 replace_with_kernel_inject 不能同时使用。

加载 Checkpoint：HuggingFace、Megatron 与 DeepSpeed 三种格式

HuggingFace 权重预加载

对 HuggingFace 训练得到的模型，最直接的方式就是前面代码里的 from_pretrained API 预先加载权重，此时 init_inference 的 checkpoint 参数可以传 None。

Megatron-LM 模型并行 checkpoint

对使用模型并行训练的 Megatron-LM 模型，需要在 JSON 配置中列出全部模型并行分片 checkpoint。以训练时 MP=2 为例：

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

DeepSpeed 训练的模型

对用 DeepSpeed 训练的模型，checkpoint 的 JSON 更简单——只需保存指向模型 checkpoint 的路径：

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

推理 MP 度可以与训练时不同。 例如：训练时完全不用 MP 的模型可以按 MP=2 做推理；训练时 MP=4 的模型也可以不切分直接推理。DeepSpeed 会在初始化阶段自动完成 checkpoint 的合并或切分。这一能力与 inference/config.py 中的 training_mp_size 配置相呼应——它用来声明 checkpoint 训练时的 MP 大小，可能与推理时的 tensor_parallel.tp_size 不同。

启动推理：deepspeed 启动器

推理脚本与训练脚本一样，使用 DeepSpeed 自带的 deepspeed 启动器拉起多卡进程，进程号会通过环境变量（如 LOCAL_RANK、WORLD_SIZE）注入：

deepspeed --num_gpus 2 inference.py

端到端实战：GPT-NEO 2.7B 多卡文本生成

下面是一份完整的客户端脚本，把 DeepSpeed-Inference 与 HuggingFace pipeline 结合，用 GPT-NEO-2.7B 做文本生成（保存为 gpt-neo-2.7b-generation.py）：

# Filename: 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)

关键点解读：

• 上面的脚本只是把 HuggingFace text-generation pipeline 内部的模型替换成了 DeepSpeed 推理引擎的产物（init_inference 返回对象的 .module）；
• 注意这里即使原始模型完全没有经过模型并行训练、checkpoint 也是单卡格式，推理阶段依然可以借助 MP 把模型按张量切分（tensor-slicing）摊到多张 GPU 上运行；
• 由于每张卡都跑同一份代码，输出只在 rank == 0 打印，避免多进程重复刷屏。

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

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

下图是脚本运行时的生成文本输出示例，你可以更换 prompt 观察该模型在不同提示下的生成效果：

[{
    'generated_text': 'DeepSpeed is a blog about the future. We will consider the future of work, the future of living, and the future of society. We will focus in particular on the evolution of living conditions for humans and animals in the Anthropocene and its repercussions'
}]

数据类型与量化模型：fp32 / fp16 / int8 与 MoQ

DeepSpeed-Inference 支持 fp32、fp16 与 int8 三种参数精度。通过 init_inference 的 dtype 参数即可设定目标数据类型，DeepSpeed 会自动挑选针对该数据类型优化的内核。从 inference/config.py 的 DtypeEnum 看，除了上述三种，还额外支持 bf16（字符串别名均可用，例如 "half"、"float"、"int8"）；字段校验器会同时接受字符串与 torch.dtype 两种写法。

对于 int8 量化模型，如果它当初是用 DeepSpeed 的 MoQ 量化方法训练出来的，那么量化施加方式的配置也必须一并传给 init_inference。这段设置包含两个信息：

• 分组数量（quantize groups）：量化整张模型用到的 scale 数量；
• MLP 部分是否带额外分组（extra grouping）：Transformer 的 MLP 部分是否用额外的分组做量化。

对应到 inference/config.py 中，quant 配置字段（仅对 int8 dtype 生效）的注释进一步说明了它如何消费 MoQ 设置：既可以传一个数值（表示量化分组数），也可以传一个元组来表示"MLP 部分存在额外分组"的语义，例如 (True, 8) 表示模型全部按 8 groups 量化、MLP 部分再额外套一层 8 grouping。关于 quantize_groups、mlp_extra_grouping 等量化参数的完整语义，可参考仓库内的 MoQ 量化教程——其中给出了 quantize_type、quantize_bits、quantize_schedule、quantize_algo、eigenvalue 等完整参数表与 GLUE 微调配置样例。

在 init_inference 中使用量化设置的示例：

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

深入源码：DeepSpeedInferenceConfig 关键参数速查

除教程正文用到的参数外，完整推理配置统一定义在 deepspeed/inference/config.py 的 DeepSpeedInferenceConfig 中。下表汇总了最具实践价值的字段及其默认值，方便你在实际工程中按需查阅：

配置字段	默认值	含义与要点
replace_with_kernel_inject（别名 kernel_inject）	False	设为 True 则自动注入高性能推理内核，适用于 Bert/GPT-2/GPT-Neo/GPT-J 等兼容模型
dtype	torch.float16	目标推理精度，支持 fp16/fp32/bf16/int8，可用字符串别名传入
tensor_parallel（别名 tp）	{}	TP 配置字典：tp_size（默认 1）、tp_grain_size（默认 64）、mpu、tp_group
enable_cuda_graph	False	为推理算子捕获 CUDA-Graph，用图回放加速（需 PyTorch ≥ 1.10）
use_triton	False	使用 Triton 内核跑推理算子（需安装 Triton，否则报错）
triton_autotune	False	开启 Triton autotune，性能更好但首次运行耗时更长
injection_policy（别名 injection_dict）	None	手动指定 {客户层: 注入策略/两个关键线性层} 的映射
triangular_masking（别名 tm）	True	控制 Transformer 层注意力分数的掩码类型，与应用场景相关
keep_module_on_host	False	加载大 checkpoint 时先留在主机内存，避免直接搬到设备导致 OOM
quant	{}	MoQ int8 量化设置，仅对 int8 dtype 生效（见上文量化小节）
checkpoint	None	DeepSpeed 兼容 checkpoint 路径，或带加载策略的 JSON 路径
base_dir	""	所有 checkpoint 文件所在的根目录，可通过 JSON 传入
max_out_tokens（别名 max_tokens）	1024	推理引擎单次可处理的输入+输出最大 token 数，长序列场景需调大
min_out_tokens（别名 min_tokens）	1	向运行时声明所需的最小输出 token 数，不足时报错而非非法内存访问
training_mp_size	1	checkpoint 训练时的 MP 大小，可与推理 tp_size 不同
moe	{}	Transformer 是否为 MoE 结构（ep_size、moe_experts、type 等子字段）
zero	{}	推理引擎可选的 ZeRO 配置
mp_size / mpu / ep_size 等	—	已弃用，请分别改用 tensor_parallel.tp_size、tensor_parallel.mpu、moe.ep_size

从配置结构可以推断：这套设计把"内核注入、张量并行、量化、MoE、checkpoint 加载"等关注点拆成了独立可组合的子配置对象，默认值经过精心选取（如 dtype 默认 fp16、max_out_tokens 默认 1024），让使用者可以从零参数起步，逐步按需打开能力。

从 v1 到 FastGen：当前版本的演进提示

本教程对应的 init_inference 机制仍可在当前仓库中正常使用，其入口与配置均保存在 deepspeed/init.py 与 deepspeed/inference 目录中。若你追求更优的生成吞吐、更新的特性与更广的模型覆盖，仓库官方推荐转向新一代方案 DeepSpeed-FastGen，其发布说明、架构与用法位于 DeepSpeed-FastGen 发布博客。理解了本文的模型并行、内核注入与 checkpoint 语义，你会更容易迁移到 FastGen 的新一代推理接口上。