【Bug已解决】[Bug]： Possible to get GPU OOM for DP/EP 解决方案-二趣网

【Bug已解决】[Bug]： Possible to get GPU OOM for DP/EP 解决方案

一、现象长什么样

在跑 MoE 大模型（DeepSeek 系列、Qwen3-MoE 等）时，同时开了数据并行（DP）和专家并行（EP），显存占用在「前几步还正常、一进真正的 forward 就炸 OOM」：

torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 2.10 GiB.

或者更隐蔽，发生在 expert 通信阶段：

OutOfMemoryError: CUDA out of memory during all-to-all dispatch (EP buffer)

几个特征：

同样的模型，只开 EP 不开 DP 能跑；只开 DP 不开 EP 也能跑；DP + EP 一起开就 OOM。
OOM 不是发生在权重加载（那是启动时），而是发生在推理进行中、tokens 被路由到 expert 时。
调小--gpu-memory-utilization能「缓解但不根治」——调太低吞吐崩，调太高还是 OOM。
在卡数多、EP degree 大时更明显。

本质：显存预算把 KV 缓存按「单副本」算满了，没给 DP 的每个副本各留一份 KV，也没给 EP 的 all-to-all 通信缓冲区预留峰值显存，于是一到 expert 路由就不够用。

二、背景

MoE 模型把 FFN 换成「很多专家」，推理时要做两件事：

专家并行（EP）：把专家分摊到不同 GPU，每个 GPU 只持有部分专家。tokens 要先通过 all-to-all 通信「dispatch」到持有对应专家的 GPU，算完再 all-to-all「combine」回来。这两个 all-to-all 各需要一块临时缓冲区装中转的 token 张量。
数据并行（DP）：同一份模型起多个数据并行副本，各处理不同请求/batch。每个 DP 副本都要独立维护自己的 KV 缓存（因为各自处理的序列不同）。

关键矛盾在显存账本上：

KV 缓存是「按 DP 副本数翻倍」的：DP=2 意味着两套 KV 缓存，显存要 ×2。但很多内存规划器（memory planner）在算 KV 预算时，默认「一份 KV 缓存」（按单副本算），于是把空闲显存几乎全部分配给了 KV —— 它不知道还会有第二个 DP 副本来要内存。
EP 的 all-to-all 缓冲区是「峰值额外开销」：dispatch/combine 时，token 张量要被重新整形并暂时放大（因为要按 expert 分组、可能填充），这块峰值内存在 forward 中途才出现，规划器如果没预留，就被 KV 缓存挤没了。

两者叠加：KV 缓存按单副本算满 + EP 通信峰值没预留 = forward 中途 OOM。

三、根因

根因是显存规划器在 DP/EP 组合下做了两个错误的「默认假设」，三层：

第一层（主因）：KV 缓存预算没乘 DP degree。规划器算「还剩多少显存给 KV」时，用的是free_mem / 1（假设单副本），但实际每个 DP rank 都要一份 KV，总 KV 显存 =per_replica_kv * dp_degree。规划器把free_mem几乎全给了「它以为的唯一一份 KV」，当第二个、第三个 DP 副本启动时各自要 KV，空闲显存早已被第一份占满 → OOM。

第二层：EP 的 all-to-all 峰值缓冲没预留 headroom。dispatch/combine 的临时缓冲在 forward 中途才分配，峰值可能达到「单 batch 所有 token × 隐藏维度 × 系数」。规划器在分配 KV 时没有扣掉这块 headroom，于是 KV 把显存吃满，all-to-all 一分配就 OOM。这解释了「为什么只发生在 tokens 被路由到 expert 时」。

第三层：规划器假设「峰值 = 权重 + KV」，忽略了「通信缓冲」这一独立项。它把显存分成「权重」「KV 缓存」两块，没单独列「通信缓冲」。在纯 TP/PP 下通信缓冲小到可忽略，但 EP 下 all-to-all 缓冲很大，这个忽略就致命了。

一句话：规划器按单副本算 KV、且漏算了 EP 的 all-to-all 峰值缓冲，导致 DP/EP 组合下 KV 缓存把显存吃满、forward 中途通信分配失败而 OOM。

四、最小可运行复现

下面用纯 Python 算一笔「显存账」，演示「KV 按单副本算满、没留 EP 缓冲」如何导致 forward 中途 OOM，不需要 GPU：

from dataclasses import dataclass @dataclass class Plan: total: int weights: int kv_per_replica: int ep_buffer_peak: int dp: int ep: int def plan_buggy(p: Plan): """有 bug 的规划器：KV 按单副本算，不留 EP 缓冲。""" free = p.total - p.weights # 错误：以为只要一份 KV kv_alloc = free reserved_ep = 0 return {"kv_alloc": kv_alloc, "reserved_ep": reserved_ep} def simulate_forward(p: Plan, plan): """模拟 forward：每个 DP 副本各要一份 KV + EP 峰值缓冲。""" kv_needed = p.kv_per_replica * p.dp ep_needed = p.ep_buffer_peak if p.ep > 1 else 0 used = p.weights + kv_needed + ep_needed return used <= p.total def main(): p = Plan(total=80, weights=20, kv_per_replica=25, ep_buffer_peak=12, dp=2, ep=4) plan = plan_buggy(p) print("规划器分配 KV:", plan["kv_alloc"], "预留 EP:", plan["reserved_ep"]) ok = simulate_forward(p, plan) print("forward 能否跑完（应 False）:", ok) # False -> OOM if __name__ == "__main__": main()

跑出来forward 能否跑完是False——因为规划器把 60GiB 全给了「一份 KV」，但实际需要25*2(DP) + 12(EP) = 62GiB给 KV+通信，加上权重 20 共 82 > 80，OOM。这就是线上的精确形状。

五、解决方案（第一层：最小直接修复）

最省事的救火：手动把--gpu-memory-utilization调低，给 EP 通信和 DP 副本留出 headroom。这不是根治，但能立刻跑起来：

# 把利用率从默认 0.9 降到 0.7，腾出 20% 显存给通信缓冲和额外 DP 副本 vllm serve <model> \ --tensor-parallel-size 1 \ --data-parallel-size 2 \ --expert-parallel-size 4 \ --gpu-memory-utilization 0.7

或者，更直接地减少 DP degree（如果业务允许）：DP=2 改成 DP=1，把省下的显存留给 EP 通信。临时规避时常用。

如果必须 DP+EP，先固定 EP 的 all-to-all 缓冲大小上限，避免峰值失控：

# 给 EP 通信缓冲设硬上限，防止它把显存吃爆 EP_BUFFER_CAP_GiB = 8

六、解决方案（第二层：结构性改进）

第一层是「手动留余量」，第二层是「让规划器自己算对」——核心是把显存账本分成四项，且 KV 预算乘 DP degree、并显式预留 EP 峰值缓冲：

from dataclasses import dataclass @dataclass class MemoryBudget: total_gib: int weights_gib: int kv_per_replica_gib: int ep_buffer_peak_gib: int dp: int = 1 ep: int = 1 headroom_gib: int = 2 # 永远留一点兜底 def plan(self) -> dict: # 1) 权重固定占用 used = self.weights_gib # 2) EP 通信缓冲（仅 EP>1 时需要） ep_buf = self.ep_buffer_peak_gib if self.ep > 1 else 0 used += ep_buf # 3) headroom used += self.headroom_gib # 4) 剩下的给 KV，但必须能装下 dp 份 remaining = self.total_gib - used kv_total_needed = self.kv_per_replica_gib * self.dp kv_alloc = min(remaining, kv_total_needed) # 关键断言：KV 预算必须至少够 dp 份，否则主动降配而非 OOM assert kv_alloc >= self.kv_per_replica_gib, ( "显存不足以支撑至少一份 KV 缓存，请降低 dp/ep 或模型规模" ) return { "weights": self.weights_gib, "ep_buffer": ep_buf, "headroom": self.headroom_gib, "kv_alloc_per_pool": kv_alloc // max(self.dp, 1), "kv_total": kv_alloc, } def advise(p: MemoryBudget) -> str: plan = p.plan() if plan["kv_total"] < p.kv_per_replica_gib * p.dp: return (f"建议降低 dp 到 {p.dp-1} 或减少 ep_buffer，" f"当前 KV 总需求 {p.kv_per_replica_gib*p.dp}GiB " f"超过可用 {plan['kv_total']}GiB") return "显存预算充足"

这样规划器在启动阶段就能算出「真正能给 KV 多少」，不会把显存吃满，也给 EP 通信留了位置。若算出来不够，它主动降级配置并给出可读建议，而不是等到 forward 中途 OOM。

七、解决方案（第三层：断言 / CI 守护）

把「KV 预算 × DP」「EP 缓冲预留」「不足即降级」固化成测试：

import pytest def test_kv_budget_multiplied_by_dp(): b = MemoryBudget(total=80, weights=20, kv_per_replica=25, ep_buffer_peak=12, dp=2, ep=4) plan = b.plan() # 总 KV 必须 >= dp 份 assert plan["kv_total"] >= 25 * 2 def test_ep_buffer_reserved_when_ep_gt_1(): b = MemoryBudget(total=80, weights=20, kv_per_replica=25, ep_buffer_peak=12, dp=1, ep=4) assert b.plan()["ep_buffer"] == 12 def test_ep_buffer_zero_when_ep_eq_1(): b = MemoryBudget(total=80, weights=20, kv_per_replica=25, ep_buffer_peak=12, dp=1, ep=1) assert b.plan()["ep_buffer"] == 0 def test_insufficient_memory_raises_not_oom(): # 显存太小，规划器应主动降级而非默默 OOM b = MemoryBudget(total=40, weights=20, kv_per_replica=25, ep_buffer_peak=12, dp=2, ep=4) with pytest.raises(AssertionError): b.plan() def test_forward_fits_after_correct_plan(): p = Plan(total=80, weights=20, kv_per_replica=25, ep_buffer_peak=12, dp=2, ep=4) bud = MemoryBudget(p.total, p.weights, p.kv_per_replica, p.ep_buffer_peak, p.dp, p.ep) plan = bud.plan() # 用规划后的 KV 总量模拟 forward，应能跑完 used = p.weights + plan["kv_total"] + plan["ep_buffer"] + plan["headroom"] assert used <= p.total

再加一个端到端回归：在模拟 DP/EP 下跑 forward，断言不触发 OOM 路径：

def test_dp_ep_forward_no_oom(): engine = make_engine(dp=2, ep=4) engine.configure_memory_budget() # 用修正后的规划器 out = engine.generate("长上下文...") # 应正常返回，不 OOM assert out is not None

八、排查清单

看 OOM 栈是否发生在all_to_all/dispatch/combine或 forward 中途，而非权重加载 → 坐实本问题。
单独关 DP（dp=1）或单独关 EP（ep=1）试，定位是不是「叠加」引发。
检查--gpu-memory-utilization：调低到 0.7 若缓解，说明是 headroom 不足。
临时救火：降gpu-memory-utilization、降 dp degree、给 EP 缓冲设上限。
长期修复：规划器把显存分四项（权重/KV×DP/EP缓冲/headroom），不足即降级。
升级 vLLM 到合了 DP/EP 显存规划修复的版本，并跑上面的 KV 预算测试。
若 OOM 只在长序列出现，说明 EP 缓冲峰值随序列长度增长，需让缓冲上限随 batch/seq 动态计算而非固定。

九、小结

DP/EP 下的 GPU OOM 不是「模型太大」，而是显存规划器按单副本算 KV、且漏算 EP 的 all-to-all 峰值缓冲，导致 KV 缓存把显存吃满、forward 中途 expert 通信分配失败。最小修复是手动调低gpu-memory-utilization或降 DP degree 留 headroom；结构性修复是把显存账本拆成「权重 / KV×DP / EP 缓冲 / headroom」四项并主动降级；最后用 pytest 把「KV 预算乘 DP」「EP 缓冲预留」「不足即降级」锁死。抓住「并行组合下显存要按最坏峰值预留」这条，TP/PP/DP/EP 任意组合的 OOM 都能用同一套路排查。

企业官网建设流程全解析