Claude opus 4.6 独立发现的 DeepEP 时序问题

原文链接: https://zhuanlan.zhihu.com/p/2013987217081644893

DeepEP Issue:https://github.com/deepseek-ai/DeepEP/issues/589

背景

前几天同事在跑实验,发现某些场景下 EP 通信方式用 DeepEP 会导致梯度变成 nan,通信方式改成 allgather 就好了。组里大哥们一起尝试复现了一下,发现了几点:

  • 可以把网络规模缩小的很小,单机就可以复现
  • 开启 CUDA_LAUNCH_BLOCKING=1 后问题消失
  • 有时候添加打印或者同步,问题又不复现
  • EP 越小越容易复现

顿时觉得头皮发麻,很显然是一种复杂的时序问题,不清楚是 megatron 的问题还是 DeepEP 的问题,也可能是在某种神秘场景,联合作用下导致的 bug。而 DeepEP 我又不怎么精通,感觉 debug 会非常困难。

vibe debug 流水账

本周三临睡前,准备试一下用 cc (Claude Code) 去定位这个问题,因为我有个稳定的复现脚本(agent 可以直接观察他是不是 nan,信号很明确),各种源码也都有(agent 能拿到所有的上下文),理论上非常适合 vibe。把复现脚本、megatron 代码、DeepEP 代码都给 claude 准备好,不过我们还没什么项目 skills,cc 的领域知识只能从源码获取,只告诉 cc 这个脚本有 bug,去研究一下为什么。yolo 一开直接睡觉。

周四感觉 vibe 情况不是非常乐观,cc 给出了几个错误的猜想都被我否定了,就开始自暴自弃,说什么问题很难定位不出来。我准备换个思路,让他先定界是 megatron 的问题还是 DeepEP 的问题,尝试把这个网络的 DeepEP 相关业务逻辑抠出来,脱离 megatron 进行复现。

周四周五公司有活动,大部分时间都在坐飞机/吃饭/喝酒/吹牛逼/听别人吹牛逼,只能每隔几个小时去看一下 cc 的 debug 情况,我和同事戏称为“收菜”。但实际上 cc 的定位进度还是堪忧,脱离 megatron 也没法复现,而且特别容易放弃,只能反复让他尝试。

复现到了周五晚上也没什么进展,准备再换个思路,不尝试脱离 megatron 复现,而是让他简化 megatron,一步步的把没用的组件剥离出去,yolo 一开继续睡觉。

周六一早,发现很牛逼,直接和我说复现出来了,我自己一跑发现确实如此,脚本的现象是这样的:

  • torch 开启 deterministic + DeepEP 有 nan
  • torch 关闭 deterministic + DeepEP 没有 nan
  • torch 开启 deterministic + DeepEP + CUDA_LAUNCH_BLOCKING 没有 nan

需要 deterministic + DeepEP + cpu 异步同时满足,非常神秘啊。yolo 一开,继续让 cc 找一下问题根因。

周六吃喝玩乐一天,晚上回家发现已经搞定了,根因如下:

1
2
3
4
5
6
7
8
9
10
11
12
// Step 1: stream_wait synchronizes comm_stream with compute_stream
// comm_stream will wait for everything *currently* enqueued on compute_stream
stream_wait(comm_stream, compute_stream);

// Step 2: torch::empty() allocates tensors AFTER the sync point
// When fill_uninitialized_memory=True, this launches a NaN-fill kernel
// on compute_stream — which comm_stream does NOT wait for
auto recv_x = torch::empty({num_recv_tokens, hidden}, x.options());

// Step 3: communication kernel runs on comm_stream, writing to recv_x
intranode::combine(..., recv_x.data_ptr(), ..., comm_stream, ...);

  • DeepEP 分配一块 empty 当输出,本来只有通信流在写数据,现在主流也会往里面写 nan,导致两个写操作并发了,制造了时序问题。这个问题的复现频率与 EP 执行速度、算子下发速度都有关系,所以看起来很复杂。
  • 本质是 torch 框架在 deterministic 的场景下破坏了约定,在 empty 的时候有写操作。但很难说是 bug 还是 feature。

思考

claude debug 成功之后,我是非常震撼的,因为这不是一个很简单的 bug,涉及到很多功能的交叉。各位 ai infra 程序员可以扪心自问一下,直接面对这个问题,大家要花几个工作日能定位出来,定位过程中心态会不会出现问题。但现在,我只是在娱乐过程中随便 prompt 了几下,claude 就全搞定了。

当然整个 vibe debug 过程没有这么顺利,如果公司没有活动,是正常的工作日,我可能对 claude 也没有这么多的耐心。这说明我们还是需要 agent 工程、需要 skills。但模型的智能已经够了,我觉得我不算很菜的程序员,但 opus 4.6 比我要强。

如何把智能充分发挥出来只是个时间问题,tipping point 已经来临了。