[YanhuiDua]

问题与修复

修复异步 all-to-all 的跨 stream 生命周期缺口：dispatch/combine 的输入由 compute stream 产生、communication stream 读取，输出则由 communication stream 分配/写入、compute stream 消费。只用 event 保证执行顺序，并不能让 caching allocator 知道另一条 stream 上尚未结束的使用。

在本次 GLM SP2 grad norm NaN 复现中，关键缺口是 _AsyncCombine.forward 输出没有登记 compute stream。该输出被后续 unpermute 的 autograd 保存供反向使用；Python/autograd 释放引用时，GPU 对旧激活的读取可能仍在排队，allocator 却已经允许 communication stream 复用其 storage。后续 _AsyncCombine.backward 的 torch.empty 在确定性模式下自动填 NaN，提前覆盖旧激活，导致 unpermute 的概率梯度 prob_grad 非有限，最终全局 grad norm NaN。这里的 NaN 填充是 PyTorch 正常的确定性保护行为，问题在于 storage 被过早复用。

本 PR：

• 为前向输入登记 communication stream，为前向和反向输出登记 compute stream，覆盖各个跨 stream 使用边界。
• 新增的输出 compute-stream 登记通过 _record_stream 作用于 AsyncCollectiveTensor（ACT）的底层 Tensor。直接调用 ACT 的 record_stream 会进入其非 view 操作分派，隐式调用 wait_tensor；底层 Tensor 登记只保护 storage 生命周期。保留现有显式 event 顺序及 dispatch 前向已有的 communication-stream 登记/等待行为，与调查中验证过的 raw_record 范围一致。
• 不引入调查用的额外 hold、分配前等待或全局同步，也不关闭确定性或未初始化内存填充。

各处登记分别保护不同输入/输出的使用，不能互相替代。本次最强直接证据指向 combine 前向输出；其他单项撤除实验通过，不足以证明可以同时删除其他登记。

调查证据

以下来自本地调查文档 e2e_38894603527_investigation_plan_2026-10-20.md 的问题 2，主要见 Q、S、U、X、Y 节。为便于审阅，在此摘录关键结果；完整本地日志及诊断脚本未纳入本 PR。

1. 保存的激活确实发生变化，并先污染 prob_grad

Q 节只撤掉完整 raw-record 补丁中 combine 前向输出的 compute-stream 登记：4 次 replay × 8 ranks 全部 norm NaN。192 个有效 forward/backward 配对中，69 个在 backward 前 checksum 改变；捕获 16 条局部链：前向激活有限 → backward 时同 ptr/storage 激活非有限，同时 incoming gradient、probs、act_grad 有限，而 prob_grad 非有限。恢复这一处登记后，192 个配对 checksum 变化及被检查的输入/梯度非有限均为 0，所有 replay/rank norm 有限。

原始未修复实现也捕获同一路径（rank4、replay2、seq34），并非仅存在于单项消融版本。此处 prob_grad 是 unpermute 对路由概率的梯度，不能据此把根因归为 topk 算法本身。

S 节独立延长 combine 输出引用保留窗口，直至完整 backward 的相关 stream 工作完成：缺失登记的基线全部 NaN；hold 组全部有限，192 个配对异常为 0。这支持生命周期解释，但该组与基线为分别启动的实验，不作为严格同权重 A/B。

2. 覆写窗口定位到分配时的 NaN 初始化

U 节 allocator trace 将捕获到的后续同址重叠分配定位到 _AsyncCombine.backward 的 torch.empty。单看通信区间覆盖旧读取，尚不能证明是 NCCL 写坏。

X 节在 torch.empty 之后、通信开始之前插入定点 reader 等待：221 次等待顺序均正确，4 × 8 norm 仍全部 NaN。保护组的 14 条局部 NaN 链中，已记录的候选通信均开始于旧读取结束之后，说明只推迟通信写入已经太晚。

同版本 PyTorch 2.9.1+cu128 的实际 GPU 微型实验中，4096 个 BF16 元素连续四次复用同一地址，fill_uninitialized_memory 开/关/开/关分别得到全 NaN/保留原值 7/全 NaN/保留 7。

3. 同一训练状态的分配前保护 ABAB 对照

Y 节沿用 8 卡、SP2、EP2 原配方，在同一进程复用同一 engine_input，四次 replay 的权重 fingerprint 和全参数 sum 签名固定。8 ranks 的 deterministic_algorithms 与 fill_uninitialized_memory 始终为 true。仅交替开启 torch.empty 之前的旧 reader 等待：

replay	分配前保护	8-rank grad norm	保存输入变化/非有限配对	prob_grad 非有限配对
0	关	全 NaN	13 / 13	14
1	开	全有限，64.475273	0 / 0	0
2	关	全 NaN	24 / 24	20
3	开	全有限，64.493820	0 / 0	0

保护两轮合计 96 个配对，保存输入 checksum 变化及输入/incoming/act_grad/prob_grad 非有限全部为 0；384 次 reader 依赖没有顺序违规。未保护组的 9 条干净局部 NaN 链均找到旧读取开始前已完成的同址重叠分配。

结合 X/Y，可将本次覆写窗口收敛到分配阶段的确定性 NaN 初始化，而非仅凭通信时间相关性归因。分配前等待因未来地址尚未知，等待范围较保守，只用于因果验证；正式修复使用 allocator 的 stream 生命周期登记。

主要本地产物：e2e_investigation_20261020/causal_saved_20261023/ 下的 combine_allocation_guard_20261023_r1_verified_results.json、combine_targeted_write_guard_20261023_r1_verified_results.json、deterministic_empty_nan_cuda_20261023.json。

验证

本次调整（最终提交 86116039，代码调整 cf6607fc）：

• CPU ACT 分派测试：2 passed；普通 Tensor 与 ACT 均执行底层登记且不调用 collective wait。将 helper 改回直接 ACT 登记的负对照按预期失败。
• AST 对照：helper 写法归一化后，dispatch/combine 四个 forward/backward 方法与调查中已验证的 raw_record 版本完全一致。
• 经指定 gpu_lock.sh 获取共享锁后执行完整测试文件：5 passed（2 项 CPU + 3 项 CUDA）；新增 combine 保存激活用例单独运行也通过（1 passed）。
• CUDA 负对照：仅删除 combine 前向输出的 compute-stream 登记，新增用例在“新 buffer 地址不得与尚在读取的旧激活相同”断言处按预期失败，实际地址相同，且新 buffer 为 NaN。测试先预热 reduction/allocator 与 backward，避免首次初始化同步干扰异步窗口。
• 完整 pre-commit 通过，含 ruff、格式检查、mypy（409 个源文件）及 pydantic；git diff --check 通过。初次沙箱中的 pydantic CUDA 初始化失败，随后在共享锁内使用训练环境重跑通过。
• 新增用例覆盖 combine 前向输出被 autograd 保存后、backward GPU 读取未结束时的 storage 保护；现有两项反向输出测试保留。collective 在单卡测试中使用本地 copy 替身，测试范围是 stream/allocator 生命周期，不等同于多卡 NCCL 集成验证。

此前调查中的修复版本已完成两轮 8 × H200、SP2 原配方 20-step 训练，均无 NaN；这是历史验证，不冒充本次调整后提交的完整训练重跑。