DeepSpeed 推理引擎初始化全解：`deepspeed.init_inference()` 入口与 `DeepSpeedInferenceConfig` 配置指南

init_inference() 是 DeepSpeed 推理（Inference）功能的统一入口：把一个普通的 torch.nn.Module 包装成高性能的 InferenceEngine，并以 DeepSpeedInferenceConfig 一个配置对象集中控制精度转换、张量并行（TP）、内核注入（Kernel Injection）、MoE、量化、CUDA Graph 等全部推理选项。本文基于 deepspeed.inference.config 的字段定义与 inference-init.rst 的官方说明展开，逐项拆解配置语义、四种合法的调用形式、嵌套子配置结构以及 InferenceEngine 内部真实的初始化流程，帮助你在不查看源码的情况下准确地把任意 HF 风格模型接入 DeepSpeed 推理管线。

从官方 API 文档出发：文档原文骨架

官方推理初始化文档（docs/code-docs/source/inference-init.rst，隶属于 API 文档导航树 index.rst 的 "Model Setup" 一节）明确了两条核心事实：

1. 入口唯一：DeepSpeed 推理的入口函数是 deepspeed.init_inference()。
2. 配置集中：DeepSpeedInferenceConfig 用于控制 InferenceEngine 初始化的所有方面；配置既可以以字典形式传给 init_inference，也可以直接以关键字参数传入。

文档给出的最小示例为：

engine = deepspeed.init_inference(model=net, config=config)

并给出了一个典型的完整配置字典示例：

config = {
    "kernel_inject": True,
    "tensor_parallel": {"tp_size": 4},
    "dtype": "fp16",
    "enable_cuda_graph": False
}

注意 "kernel_inject" 在这里通过 pydantic 字段别名（alias）映射到内部的 replace_with_kernel_inject，而 "tensor_parallel" 也可以简写为 "tp"——这些别名细节将在下文逐一说明。返回的 engine 是一个 InferenceEngine，它会替换模型内部的关键算子并接管前向计算，因此你原有的调用方式（例如 HF pipeline 或模型自带的 generate）在包装后仍可继续使用。

入口函数 init_inference：签名、四种调用形式与合并规则

deepspeed.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：可选，可以是配置字典，也可以是指向 JSON 配置文件的路径字符串；
• **kwargs：可选，所有 DeepSpeedInferenceConfig 字段均可作为关键字参数直接传入。

根据 init.py 的 docstring 与实现，以下四种情况全部合法：

Case 1：不传任何配置，使用全默认值。

generator.model = deepspeed.init_inference(generator.model)
string = generator("DeepSpeed is")
print(string)

Case 2：只传 config 字典。

config = {"tensor_parallel": {"tp_size": 4}, "dtype": "fp16", "kernel_inject": True}
generator.model = deepspeed.init_inference(generator.model, config=config)

Case 3：只传关键字参数（不传 config）。

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

Case 4：config 与 kwargs 同时提供（合并语义）。

generator.model = deepspeed.init_inference(generator.model,
                                           config={"dtype": torch.half},
                                           replace_with_kernel_inject=True)

config 与 kwargs 的合并规则（源码级）

config 与 kwargs 并不是简单叠加，合并逻辑 顺序如下：

1. config 为 None 时视为空字典；为字符串时按 JSON 文件读取（json.load）；为其他类型则抛出 ValueError；
2. 取 config_dict 与 kwargs 的交集键 overlap_keys；
3. 对交集键逐一比对：若同一个键在两边取值不同，立即抛出 ValueError(f"Conflicting argument '{key}' ...")，而不是静默覆盖；
4. 取值相同的重叠键与各自独有的键最终通过 config_dict.update(kwargs) 合并。

也就是说，代码实现中"kwargs 优先"只在两边取值一致（无冲突）时才成立；一旦冲突会直接报错。这一点被单元测试 tests/unit/inference/test_inference_config.py 显式覆盖：test_overlap_kwargs 验证同值重叠正常通过，test_overlap_kwargs_conflict 验证异值重叠抛出 ValueError，test_kwargs_and_config 验证互补键合并。另有一个值得注意的测试 test_json_config（同一文件）证明 config 可以直接传 JSON 文件路径。

合并完成后，函数将结果实例化为 DeepSpeedInferenceConfig(**config_dict)，再构造 InferenceEngine(model, config=ds_inference_config) 并返回。

顶层配置类 DeepSpeedInferenceConfig：全部字段与别名速查

整个推理初始化由 DeepSpeedInferenceConfig 统一建模，它继承 DeepSpeedConfigModel（pydantic 模型），因此支持别名、校验器与默认值自动填充。以下按用途分组整理字段（默认值与别名均取自源码）：

字段	默认值	别名	说明
replace_with_kernel_inject	False	kernel_inject	置 True 时对 BERT、GPT2、GPT-Neo、GPT-J 等模型注入优化后的推理内核；否则需通过 injection_dict 指定注入策略
dtype	torch.float16	—	目标数据类型，模型会被转换到该类型；文档声明支持 torch.half、torch.int8、torch.float，DtypeEnum 还允许 bf16
tensor_parallel	{}	tp	张量并行配置，接受 DeepSpeedTPConfig 字典
enable_cuda_graph	False	—	是否对推理算子捕获 CUDA-Graph 以用图回放方式加速
use_triton	False	—	是否用 Triton 内核执行推理算子；要求已安装 Triton，否则校验器抛 ValueError
triton_autotune	False	—	是否开启 Triton autotune；开启后性能更好但首次运行变慢
zero	{}	—	与推理引擎搭配的 ZeRO 配置，接受 DeepSpeedZeroConfig 字典
triangular_masking	True	tm	控制 Transformer 层注意力分数的掩码类型，与应用场景相关
moe	{}	—	是否/如何按 MoE Transformer 处理，接受布尔值或 DeepSpeedMoEConfig 字典
keep_module_on_host	False	—	加载 checkpoint 时是否保留在主机内存不搬到设备（适合超大模型，避免 OOM）
quant	{}	—	量化设置，接受 QuantizationConfig 字典；注意仅对 int8 dtype 生效
checkpoint	None	—	DeepSpeed 兼容 checkpoint 路径，或带加载策略的 JSON 路径
base_dir	""	—	所有 checkpoint 文件所在根目录
set_empty_params	False	—	推理模块创建时使用空 Tensor 还是真实 Tensor
save_mp_checkpoint_path	None	—	保存已加载模型的路径，用于调整并行度以缓解加载开销；不传则不保存
checkpoint_config	{}	ckpt_config	checkpoint 相关配置，接受 InferenceCheckpointConfig 字典
return_tuple	True	—	Transformer 层返回 tuple 还是 Tensor
training_mp_size	1	—	若加载 checkpoint，这是其训练时的模型并行度，可不同于推理时并行度
injection_policy	None	injection_dict	客户 nn.Module 到注入策略的映射字典，例如 {BertLayer: deepspeed.inference.HFBertLayerPolicy}
injection_policy_tuple	None	—	内部使用的策略元组（一般由 injection_dict 推导）
max_out_tokens	1024	max_tokens	引擎一次最多可处理的 token 数（含输入与输出），应按用例适当调大
min_out_tokens	1	min_tokens	期望生成的最少 token 数；不满足时运行时给出带内存压力的错误提示，而非段错误或脏输出
transposed_mode	False	transposed_mode	转置模式开关

废弃字段（兼容保留，使用 json_schema_extra 标注迁移方向）：

• mp_size（默认 1）→ 迁移到 tensor_parallel.tp_size；
• mpu → 迁移到 tensor_parallel.mpu；
• ep_size / ep_group(别名 expert_group) / ep_mp_group(别名 expert_mp_group) / moe_experts / moe_type → 全部迁移到 moe 子配置；
• replace_method（默认 "auto"）→ 已标记 deprecated，提示从调用中移除。

dtype 的灵活写法

dtype 字段带一个 mode="before" 的 校验器：字符串会被 DtypeEnum.from_str 归一化到对应 torch.dtype。DtypeEnum（deepspeed/inference/config.py）接受的字符串包括：

• fp16 / float16 / half / torch.float16 → torch.float16
• fp32 / float32 / float / torch.float32 → torch.float32
• bf16 / bfloat16 / bfloat → torch.bfloat16
• int8 → torch.int8

因此文档示例中的 "dtype": "fp16" 会被自动换算为 torch.float16。InferenceEngine 初始化时还会校验 config.dtype 是否在当前加速器 supported_dtypes() 列表中（见 engine.py），不支持的组合直接抛错。

四个嵌套子配置：TP、MoE、Quantization 与 Checkpoint

原文档通过 autodoc 导出了四个子配置类，它们是 DeepSpeedInferenceConfig 中 tensor_parallel、moe、quant、checkpoint_config 四个字段的取值骨架。

DeepSpeedTPConfig（字段 tensor_parallel / 别名 tp）

DeepSpeedTPConfig 定义于 deepspeed/inference/config.py：

字段	默认值	说明
enabled	True	张量并行总开关
tp_size	1	用张量并行切分模型到多少设备；1 表示不并行
tp_grain_size	64	期望的 MLP/lm_head TP 切分粒度；DNN 库偏好 2 的幂粒度，默认取 64
mpu	None	实现 get_{model,data}_parallel_{rank,group,world_size}() 的模型并行单元对象
tp_group	None	可直接指定的张量并行通信组（内部优先使用 mpu 或自行创建）

DeepSpeedMoEConfig（字段 moe）

DeepSpeedMoEConfig 用于 MoE Transformer：

字段	默认值	说明
enabled	True	MoE 总开关
ep_size	1	专家并行（expert-parallel）大小，用于把专家切分到 EP 组内的各 GPU
moe_experts（别名 num_experts）	[1]	单个 MoE 层的全局专家数
type	standard	MoE 层类型：standard 或 residual
ep_group（别名 expert_group）、ep_mp_group	None	可传入预创建的通信组

顶层字段 moe 的类型是 Union[bool, DeepSpeedMoEConfig]：直接传布尔值时，moe_backward_compat 校验器 会将其自动包装为 DeepSpeedMoEConfig(enabled=value)，对应测试见 test_moe_backward_compat_bool。

QuantizationConfig（字段 quant）

QuantizationConfig 及其子类构成层级量化配置。官方注释提醒：该功能只对 int8 dtype 生效。其结构为：

QuantizationConfig
├── enabled: bool = True
├── activation: ActivationQuantConfig   # enabled=True
├── weight:    WeightQuantConfig        # enabled=True
│     ├── quantized_initialization: Dict = {}
│     └── post_init_quant: Dict = {}
└── qkv:       QKVQuantConfig           # enabled=True

其中 BaseQuantConfig（config.py）提供公共字段：num_bits（默认 8）、q_type（symmetric/asymmetric）、q_groups（默认 1）。这里对应文档中"MoQ"（Model Quantization）能力的配置入口。

InferenceCheckpointConfig（字段 checkpoint_config / 别名 ckpt_config）

InferenceCheckpointConfig 只有三个可选路径字段，用于声明 checkpoint 存放位置：

class InferenceCheckpointConfig(DeepSpeedConfigModel):
    checkpoint_dir: Optional[str] = None
    save_mp_checkpoint_path: Optional[str] = None
    base_dir: Optional[str] = None

可复制的完整配置示例

把原文档的示例扩充为三种等价写法。

写法一：全量字典 + JSON 文件

字典写法（含 MoE、量化与 token 长度约束的进阶示例）：

import torch
import deepspeed

config = {
    "kernel_inject": True,                       # 等价 replace_with_kernel_inject=True
    "tensor_parallel": {"tp_size": 4},           # 等价 "tp": {"tp_size": 4}
    "dtype": "fp16",                             # 字符串写法会被归一化为 torch.float16
    "enable_cuda_graph": False,
    "return_tuple": True,
    "max_out_tokens": 2048,                      # 按生成任务增大上限
    "min_out_tokens": 1,
    # MoE 模型才需要开启：
    # "moe": {"enabled": True, "ep_size": 2, "moe_experts": 8, "type": "standard"},
    # int8 量化才需要开启：
    # "quant": {"enabled": True, "weight": {"num_bits": 8, "q_type": "symmetric"}},
    # "dtype": "int8",
}
engine = deepspeed.init_inference(model=net, config=config)

由于 init_inference 支持字符串路径，同一份配置也可写成 JSON 后从文件加载（对应测试 test_json_config）：

engine = deepspeed.init_inference(model=net, config="/path/to/inference_config.json")

写法二：关键字参数（kwargs）

engine = deepspeed.init_inference(net,
                                  kernel_inject=True,
                                  tensor_parallel={"tp_size": 4},
                                  dtype=torch.float16,
                                  enable_cuda_graph=False)

写法三：与 HF pipeline 组合（测试中的真实用法）

仓库集成测试 tests/unit/inference/test_inference.py 演示了最贴近实战的姿势——先构建 HF pipeline，再把 pipe.model 包进 DeepSpeed：

from transformers import pipeline

pipe = pipeline("text-generation", model="openai-community/gpt2", device="cpu")
pipe.model = deepspeed.init_inference(pipe.model,
                                      mp_size=self.world_size,   # 旧写法；新写法用 tensor_parallel={"tp_size": ...}
                                      dtype=dtype,
                                      replace_with_kernel_inject=True)
pipe.device = torch.device("cuda")
out = pipe("DeepSpeed is", do_sample=False, max_length=20)

包装后 pipe.model 仍保留 generate 能力：InferenceEngine 在检测到模型带 generate 方法时会挂接自己的 _generate（见 engine.py）。

源码级原理：InferenceEngine 初始化到底做了什么

InferenceEngine 定义于 deepspeed/inference/engine.py，其 __init__（engine.py）按以下顺序完成包装，理解这一流程有助于判断配置错误发生在哪一步：

1. 精度转换：config.dtype 非空时调用 _convert_to_dtype 把整个模型转到目标类型（engine.py）。
2. 张量并行组创建：若提供了 mpu，则以 mpu 的模型并行组为准设置 tp_size；否则当 tp_size > 1 时调用 init_distributed() 并按 range(tp_size) 创建 mp_group（engine.py 与 _create_model_parallel_group）。mp_group 会被缓存为类级变量 inference_mp_group 以便复用。
3. 三种并行/注入模式互斥（engine.py 的注释直接写明了"我们只支持三种模式"）：
   • 模式一：用户指定注入策略——提供了 injection_dict（即 injection_policy）时，按映射对指定 client_module 套用注入策略；同时禁止再开 replace_with_kernel_inject（assert 拒绝两者同用）。
   • 模式二：内核注入——replace_with_kernel_inject=True 且无 injection_dict 时，调用 _apply_injection_policy 走内置内核注入路径。
   • 模式三：自动张量并行（AutoTP）——既无注入字典也未开内核注入，但 tp_size > 1 时，用 AutoTP.tp_parser(model) 自动解析模型结构并逐层注入。
4. 设备搬运：支持 meta 设备模型（先 to_empty）；否则默认 self.module.to(device)；只有 keep_module_on_host=True 时才把参数留在主机端以预留量化等二次处理空间（engine.py）。
5. MoE 与 EP 组：当检测到模型含 MoE 层且 world_size > 1 时，按 config.moe.ep_size 创建专家并行组与专家-模型并行组（engine.py 及 _create_ep_parallel_group）。
6. CUDA-Graph 约束校验：使用 CUDA Graph 需要 torch ≥ 1.10（engine.py）；模型并行（tp_size > 1）时明确不支持 CUDA Graph（engine.py）。回归测试 test_cuda_graph_with_kernel_inject_raises 还证实：return_tuple=False + enable_cuda_graph=True + 内核注入的组合会抛出 ValueError("enable_cuda_graph is not supported ...")（对应 DeepSpeed 的历史 issue #8330）。

另外，InferenceEngine 也承担了针对特定架构的兼容处理：例如对 BLOOM 移除 HF 侧 _prepare_attn_mask 以适配内核注入（engine.py），以及在 tp_size > 1 时为 BLOOM/MPT 等模型重建 alibi / attention bias 的按头切分辅助逻辑（build_alibi_tensor / build_attn_bias）。这些细节从源码结构看主要是为了让"头数"随 TP 组正确分片，避免多头注意力在并行切分后出错。

常见问题速查与调试建议

结合字段默认值、源码断言与测试约束，可总结出以下排查要点：

• 想跑多卡但报并行组错误：检查是否把旧字段 mp_size 与新的 tensor_parallel.tp_size 混用（旧字段仍可用但已标记废弃，且与 tp 配置语义并存时容易困惑）；从源码结构看，二者最终都会汇聚到 tp_size 逻辑上。
• CUDA Graph 与模型并行二选一：tp_size > 1 与 enable_cuda_graph=True 互斥，开启前请确认目标架构允许。
• int8 量化不生效：确认 dtype 必须为 int8 且 quant.enabled=True；量化参数层级为 quant.weight / quant.activation / quant.qkv。
• 生成被截断：检查 max_out_tokens（默认 1024，含输入输出总 token 数）是否小于任务实际需要，必要时调大。
• 超大模型加载 OOM：使用 keep_module_on_host=True 让 checkpoint 暂留主机，先量化再搬上设备。
• config 与 kwargs 冲突：init_inference 会针对同一键的冲突取值直接抛 ValueError，而非静默覆盖；请保证两者取值一致。

关于各字段的完整语义与 autodoc 注释，可直接查阅 deepspeed/inference/config.py 与官方 API 文档 inference-init.rst；想了解 InferenceEngine 包装后提供的前向与 profile 能力，可继续阅读同目录下的 inference-engine 相关文档与 engine.py 实现。