Prhub
← 返回仓库详情

#5591 [fsdp] fix: pass dp_group to prepare_dynamic_batch to fix CUDA deadlock

verl-project/verl · 作者 JenniferWang · 合并时间 2026-03-26 16:06

分析状态 已生成
文件变更 0提交数 1 · 评论 0
代码增减 +0 / -0
fsdp worker misc

执行摘要

修复 FSDP 训练中动态批处理导致的 CUDA 死锁，通过传递 dp_group 参数确保 micro-batch 计数同步。

根据PR body，根因是commit f5c34bb在verl/utils/seqlen_balancing.py添加了dp_group is not None guard（第391行），导致FSDP actor的prepare_dynamic_batch调用未传递dp_group时跳过all_reduce。在动态批处理下，不同rank由于序列长度分布不同计算不同micro-batch数量，进而导致FSDP collectives（AllGather/ReduceScatter）死锁。

该PR值得精读以理解FSDP同步机制和动态批处理的实现细节。关键设计决策是使用WORLD作为dp_group确保所有rank同步，工程师可以学习如何避免分布式训练中的死锁问题。

讨论亮点

Review中仅有gemini-code-assist[bot]的评论，肯定了修复的有效性。评论指出：'The fix, which involves passing torch.distributed.group.WORLD as the dp_group to prepare_dynamic_batch in both compute_log_prob and update_policy, is a robust solution.' 无争议或深入讨论，修复直接被wuxibin89批准。

实现拆解

实现集中于单个文件verl/workers/actor/dp_actor.py。关键改动是在两个函数中修改prepare_dynamic_batch调用：

  1. 在compute_log_prob函数（第468行附近）添加dp_group=torch.distributed.group.WORLD。
  2. 在update_policy函数（第560行附近）添加相同参数。这确保了在use_dynamic_bsz=True时，所有数据并行rank通过all_reduce(MAX)同步micro-batch计数。
文件 模块 状态 重要度
verl/workers/actor/dp_actor.py workers/actor modified 6.0

分析完成后，这里会展示 LLM 生成的相对完整源码片段和详细注释。

关键符号

compute_log_prob update_policy

评论区精华

修复有效性确认 正确性

gemini-code-assist[bot] 评论修复是健壮的解决方案，解决了根因问题。

结论：修复正确且有效，无争议。 · 已解决

风险与影响

风险较低：

  • 正确性风险：使用WORLD作为dp_group在所有FSDP并行配置（纯FSDP、HSDP、Ulysses SP）下均正确，PR body已验证。
  • 回归风险：变更仅添加参数，不影响现有逻辑；但需确保其他调用prepare_dynamic_batch的地方也正确处理dp_group。
  • 测试覆盖：PR提供了在2-node EKS集群上的测试结果，显示修复后4/4运行干净，但没有单元测试覆盖。

影响范围：使用FSDP和动态批处理的训练场景。影响程度：高，修复了导致100%再现死锁的bug，提升训练稳定性和可靠性。对用户：确保动态批处理功能正常，避免训练中断。

动态批处理同步风险 缺少单元测试

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接，后续同步到相关引用后会出现在这里。

完整报告

Markdown PDF

执行摘要

本PR修复了FSDP训练中因动态批处理micro-batch计数不同步导致的CUDA死锁问题，通过向prepare_dynamic_batch函数传递dp_group参数确保所有数据并行rank同步，提升训练稳定性。

功能与动机

动机源于commit f5c34bb在verl/utils/seqlen_balancing.py中添加了dp_group is not None guard，导致FSDP actor调用prepare_dynamic_batch时未传递dp_group参数，跳过all_reduce操作。在动态批处理场景下，不同rank因序列长度分布差异计算不同micro-batch数量，进而引发FSDP collectives死锁。PR body指出："不同序列长度分布导致不同rank计算不同微批次计数...FSDP collectives (AllGather/ReduceScatter)死锁"。

实现拆解

实现集中于单个文件verl/workers/actor/dp_actor.py：

  • 在compute_log_prob函数第468行附近，修改prepare_dynamic_batch调用，添加dp_group=torch.distributed.group.WORLD。
  • 在update_policy函数第560行附近，进行相同修改。

关键代码片段：

micro_batches, batch_idx_list = prepare_dynamic_batch(
    data, max_token_len=max_token_len, dp_group=torch.distributed.group.WORLD
)

这确保所有rank通过all_reduce(MAX)同步micro-batch计数，防止死锁。

评论区精华

Review讨论简单，仅有gemini-code-assist[bot]的评论：

"The fix, which involves passing torch.distributed.group.WORLD as the dp_group to prepare_dynamic_batch in both compute_log_prob and update_policy, is a robust solution."

评论肯定了修复的有效性，无争议点。wuxibin89直接批准。

风险与影响

  • 风险：变更引入dp_group参数，使用WORLD在所有FSDP并行配置下正确，但需确保其他调用prepare_dynamic_batch的地方也正确处理该参数。PR提供了在2-node EKS集群的测试验证，但缺少单元测试覆盖。
  • 影响：修复了FSDP动态批处理训练中的死锁bug，影响所有使用该配置的用户，提升训练可靠性和性能。测试显示修复后4/4运行成功，而死锁前100%再现。

关联脉络

本PR是#5451的FSDP对应版本，#5451修复了megatron workers中的相同bug。这表明动态批处理同步问题在多个并行策略中普遍存在，需要跨模块统一处理。近期历史PR如#5604涉及FSDP workers重构，但本PR专注于具体bugfix。

参与讨论