【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 副本数翻倍」的DP2 意味着两套 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不需要 GPUfrom 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(total80, weights20, kv_per_replica25, ep_buffer_peak12, dp2, ep4) 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 80OOM。这就是线上的精确形状。五、解决方案第一层最小直接修复最省事的救火手动把--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如果业务允许DP2 改成 DP1把省下的显存留给 EP 通信。临时规避时常用。如果必须 DPEP先固定 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 通信缓冲仅 EP1 时需要 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(total80, weights20, kv_per_replica25, ep_buffer_peak12, dp2, ep4) plan b.plan() # 总 KV 必须 dp 份 assert plan[kv_total] 25 * 2 def test_ep_buffer_reserved_when_ep_gt_1(): b MemoryBudget(total80, weights20, kv_per_replica25, ep_buffer_peak12, dp1, ep4) assert b.plan()[ep_buffer] 12 def test_ep_buffer_zero_when_ep_eq_1(): b MemoryBudget(total80, weights20, kv_per_replica25, ep_buffer_peak12, dp1, ep1) assert b.plan()[ep_buffer] 0 def test_insufficient_memory_raises_not_oom(): # 显存太小规划器应主动降级而非默默 OOM b MemoryBudget(total40, weights20, kv_per_replica25, ep_buffer_peak12, dp2, ep4) with pytest.raises(AssertionError): b.plan() def test_forward_fits_after_correct_plan(): p Plan(total80, weights20, kv_per_replica25, ep_buffer_peak12, dp2, ep4) 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(dp2, ep4) engine.configure_memory_budget() # 用修正后的规划器 out engine.generate(长上下文...) # 应正常返回不 OOM assert out is not None八、排查清单看 OOM 栈是否发生在all_to_all/dispatch/combine或 forward 中途而非权重加载 → 坐实本问题。单独关 DPdp1或单独关 EPep1试定位是不是「叠加」引发。检查--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 都能用同一套路排查。