分布式推理账：模型放不下时，状态放哪里

推理和训练不同：训练只要吞吐，推理还要延迟和状态管理。模型单卡放不下时，怎么切、状态放哪里、扩缩容怎么办，每一步都有显存与带宽的账。本文把分布式推理的账算清楚，方法框架参考开源书稿《深入理解 AI Infra》（Apache-2.0），公式与算例为本站自行推导整理。

1. 显存账：单卡放不下 → 张量并行

张量并行（TP）把权重沿隐藏维切开，N 卡各存 1/N：

$$Mem\_per\_card = \frac{W}{TP} + \frac{KV}{TP} + 激活余量$$

70B FP16（140 GB 权重）+ 32K 上下文 KV（≈5 GB）：TP4 后每卡 ≈ 36 GB，从"放不下"变成"可运行"——这是 70B 在 48G 卡上能跑起来的量化原因。

2. 通信账：TP 的隐藏成本

TP 每个算子前后都要 All-Reduce 同步（前后两次）：

$$Comm\_per\_token \approx 2 \times \frac{TP-1}{TP} \times H_{hidden} \times bytes \times 层数(约)$$

每 token 的 TP 通信量与隐藏维度 × 层数成正比，与批大小无关。所以：

• TP 适合显存不够的场景，不是"更快"的手段；
• 批小时 TP 通信占比高，批大时数据并行（DP）更划算；
• 判断：先问"显存够不够"，不够才上 TP；够的话 DP 优先。

方案	显存/卡	每 token 通信	适用
DP（完整副本）	W+KV 全量	无（仅请求路由）	显存够、高并发
TP2	一半	1×H×bytes×层	显存临界
TP4	1/4	1.5×H×bytes×层	大模型、显存不足
TP8	1/8	1.75×H×bytes×层	超大模型

3. 状态账：KV 放哪里

KV 缓存是"会话状态"，分布式推理的核心问题是状态放置：

放置	优点	代价
与权重同卡（TP 卡）	无额外传输	占显存，限制并发
独立 KV 节点/卸载	权重层利用率高	每次生成要跨节点取 KV
prefill/decode 分离（PD）	两阶段各自优化	状态迁移（prefix 缓存命中时省）

PD 分离的账：prefill 是算力受限（一次性算长输入）、decode 是带宽受限（逐 token 读权重），两者最优配置不同。把 prefill 与 decode 拆到不同机器组，各按自己的瓶颈选卡型，是规模化推理的标准打法（见 推理引擎优化 与 vLLM 生产部署）。

4. 扩缩容与恢复账

• 无状态部分（权重）可随时加副本：加 N 个 DP 副本 → 吞吐 ≈ ×N，前提是请求能打散；
• 有状态部分（KV）扩容难：正在生成的会话无法随意迁移，prefix cache 命中可"接续"；
• 恢复：单卡故障 → 该卡上的会话断，无状态副本自动补位，KV 缓存按快照重建。

5. 可复算 Python 函数

def tp_account(n_params_b, kv_gb, tp, bytes_per_param=2, h_dim=8192, layers=80):
    """TP 推理的显存/卡 与 每 token 通信量（MB）"""
    mem = n_params_b * 1e9 * bytes_per_param / tp + kv_gb / tp
    comm_mb = 2 * (tp - 1) / tp * h_dim * layers * 2 / 1e6
    return mem, comm_mb

tp_account(n_params_b=70, kv_gb=5.2, tp=4)
# (36.3 GB/卡, 2.4 MB/token)

6. 交付清单

1. 先算显存账决定 DP/TP：显存够则 DP 优先，不够再上 TP；
2. TP 选择从 TP2 起步实测通信占比，别一步 TP8（通信账随 TP 涨）；
3. KV 状态明确放置策略：同卡 / 卸载 / PD 分离按并发与延迟目标选；
4. 扩容只承诺无状态部分（DP 副本），有状态部分按 prefix 命中率估算收益；
5. 单卡故障恢复演练一次，确认会话断连与补位行为符合预期。

---

本文为本站自行推导整理，方法框架参考开源书稿《深入理解 AI Infra》（Apache-2.0）；具体版本号与实测数据以现场为准。