DeepSpeed-Inference 与 DeepSpeed-MII

10.1 DS 推理生态的三层

实战中 vLLM / SGLang 后来居上，DS 推理生态的市场份额不如训练那么强势。但在两个场景下仍然有竞争力：

• 已经训完模型，想用 "训练 + 推理同一框架"，减少格式转换；
• 对 NVIDIA 之外硬件（AMD / NPU）有需求，DS 支持广。

10.2 DeepSpeed-Inference v1：单卡 / 多卡 TP 推理

最小示例

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import deepspeed

tok = AutoTokenizer.from_pretrained("Qwen/Qwen3-4B-Instruct-2507")
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-4B-Instruct-2507", torch_dtype=torch.bfloat16
)

# ★ 注入推理引擎
engine = deepspeed.init_inference(
    model,
    tensor_parallel={"tp_size": 2},        # 用 2 张卡做 TP
    dtype=torch.bfloat16,
    replace_with_kernel_inject=True,       # 替换 fused kernel
    max_out_tokens=2048,
)

prompt = "你好，介绍一下 DeepSpeed"
inputs = tok(prompt, return_tensors="pt").to(engine.module.device)
out = engine.module.generate(**inputs, max_new_tokens=128)
print(tok.decode(out[0], skip_special_tokens=True))

启动方式：

deepspeed --num_gpus 2 infer.py

init_inference() 关键参数

10.3 DeepSpeed-Inference v2 / FastGen

v1 偏静态批，v2（FastGen） 加入了：

• continuous batching（同 vLLM）；
• Dynamic SplitFuse —— 把 prefill 和 decode 在同一 batch 里混合，提高利用率；
• FlashAttention v2 集成；
• 更细的 KV cache 管理。

用法：

from deepspeed.inference.v2 import build_hf_engine

engine = build_hf_engine(
    path="Qwen/Qwen3-4B-Instruct-2507",
    engine_config={
        "tensor_parallel": {"tp_size": 2},
        "max_ragged_sequence_count": 256,
        "max_ragged_batch_size": 8192,
    },
)
# 之后通过 generator API 提交请求

FastGen 在同样硬件上吞吐通常和 vLLM 相当，但调用 API 风格不同；选哪个看团队习惯。

10.4 DeepSpeed-MII：开箱即用的服务

MII = "Model Implementations for Inference"。它包了 v1/v2，让你一行命令起 OpenAI 兼容 API：

pip install deepspeed-mii

# 启动 RESTful 服务
python -c "
import mii
client = mii.serve('Qwen/Qwen3-4B-Instruct-2507', tensor_parallel=2)
"

# 客户端调用
curl http://localhost:28080/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{"model":"default","messages":[{"role":"user","content":"你好"}]}'

MII 还支持 persistent 部署（不退出，挂着等请求）和 non-persistent（短任务进出）。在 LLaMA-Factory ch12 提到的"训完直接部署"流程，DS 用户的一条路就是 LLaMA-Factory 出 merged HF model → MII 部署。

10.5 训练 → 推理无缝衔接

关键链路：训练时的 ZeRO-3 切分 → 保存时通过 stage3_gather_16bit_weights_on_model_save 合回完整 fp16 → 推理时再用 DS-Inference / MII 加载。这就是 DS 一条龙生态的最大卖点。

10.6 选型对照表

10.7 重要注意

• 训练用的 model_engine 和推理用的 inference_engine 不通用：训练 checkpoint 必须先 gather 成 HF 格式，再加载到推理引擎；
• kernel inject 不是所有模型都支持：DS-Inference 维护了一份模型白名单，新架构可能需要手写适配，落到 deepspeed/module_inject/ 下；
• 量化推理的精度损失要测：DS-Inference int8 在某些任务上掉点明显，先小规模 benchmark；
• MII 默认监听 28080 / 50050，端口和 vLLM/HF TGI 不冲突，可以同机并存。

10.9 InferenceEngine 的三种 injection 路径

翻开 deepspeed/inference/engine.py 的 InferenceEngine.__init__（line 45–184）能看到一段判分支：到底用哪条 injection 路径，由三个 config 字段决定。

判定优先级：1 > 2 > 3。模型不在 kernel_inject 白名单（如某些新出的开源模型）时，会自动落到 3（AutoTP）；功能上能跑，但比 2 慢，因为没有 fused kernel。

诊断："我开了 replace_with_kernel_inject=True 但速度跟没开一样" → 说明你的模型类没在 replace_policies 里。打开 DS log，看启动期是否打印 replace_module: True 或 AutoTP: True 来确认走了哪条路径。

10.10 v2 / FastGen 的 InferenceEngineV2 API

v1 的 generate() 是静态批，v2 设计为 "持续接收请求 / 持续吐结果" 的服务端 API。源码在 deepspeed/inference/v2/engine_v2.py。

两个核心数据结构：

• DSStateManager：持续维护"哪条 sequence 占了哪些 KV-cache block"。free_blocks 属性能查剩多少空闲 block；
• RaggedBatchWrapper：把不等长的 sequences 拼成一个 GPU friendly 的 tensor，自带 metadata（offset、cu_seqlens）。

FastGen 的 "Dynamic SplitFuse" 就是在 put() 之前做的：把当前队列里的 prefill 请求和 decode 请求合并成一个 batch（prefill 算很多 token、decode 算 1 个 token），让 GPU 一次 forward 同时服务两类负载。

10.11 module_inject policy 怎么写

想给一个 kernel_inject 没支持的新模型加 v1 推理加速？需要在 deepspeed/module_inject/ 写一个 TransformerPolicy 子类。最少实现 3 个方法：

# 简化骨架
class MyModelPolicy(TransformerPolicy):
    def attention(self):
        # 返回 QKV / output 权重的张量切割信息：
        # (qkvw, qkvb, dense_w, dense_b, head_dim, ...)
        ...
    def mlp(self):
        # 返回 MLP 的 4 个张量：intermediate / output 权重和 bias
        ...
    def layernorm(self):
        # 返回 attention / mlp 前的 layernorm 权重
        ...

然后在 replace_policy.py 里把 HFMyModelLayer → MyModelPolicy 注册进 replace_policies dict。一次性活——做完后用户写 replace_with_kernel_inject=True 就能享受 fused kernel。这是 DS 推理生态对接新模型的标准 PR 路径，git log deepspeed/module_inject/ 能看到一长串以 "Add inference policy for X" 命名的 commit。

10.12 MII 实际跑的是 v1 还是 v2？

MII 本身是个 "facade"——它根据传入的 engine_config 决定起 v1 还是 v2：

启动后用 curl http://localhost:28080/v1/models 看 metadata 里有没有 "engine": "v2" 字样能确认。

10.13 这章带走什么

• 知道 DS 推理有 v1 / v2 / MII 三层；
• 会用 deepspeed.init_inference() 跑一次最小推理；
• 能看懂三种 injection 路径（user policy / kernel inject / AutoTP）的优先级；
• v2 service 化时知道关键 API 是 put / query / can_schedule，背后是 DSStateManager 在管 KV-cache；
• 清楚训练 checkpoint → HF 格式 → 推理引擎的转换路径；
• 知道在团队选型时，vLLM / SGLang 仍然是吞吐冠军，DS 推理的优势在"全栈一致"。