Skip to content

fix: fix reducer behavior when grad accumulation is on#183

Merged
kilinchange merged 2 commits into
masterfrom
fix/reducer_grad_accu
Jul 3, 2026
Merged

fix: fix reducer behavior when grad accumulation is on#183
kilinchange merged 2 commits into
masterfrom
fix/reducer_grad_accu

Conversation

@Chamberlain0w0

Copy link
Copy Markdown
Contributor

修复了梯度累积情况下 ddp 分桶的 reducer rebuild buckets 后可能导致的错误行为。

背景

  1. 第一个 step backward 的时候会记录运行时 tensor 实际 ready 顺序,然后根据这个顺序重新给 tensor 分桶,调用 RebuildBuckets(),此行为可能导致后一轮 tensor->grad() 需要重新绑定到新的 bucket view;

  2. 每一轮 tensor->grad() 重新绑定到新的 bucket view,目前的逻辑是:

    • 如果 grad 为空(ZeroGrad(set_to_none=True) 的结果)
    • 或者 grad 的指针和 bucket view 的指针不是同一个(第一点中 rebuild buckets 的结果)
      • 这种情况会报一个 warning(当时想的是用户可能会手动改 tensor->grad() 导致的)

    则把 grad 重新 set 为 bucket view,并且标志下一次 overwrite 而不是直接 accumulate(因为 bucket view 仍然可能是上一轮的脏数据。

  3. 梯度累积可能会执行 N 次 forward/backward 后,才执行一次 ZeroGrad()

错误成因

场景:第一轮后进行了 RebuildBuckets(),导致 grad 的指针和全新 bucket view 的指针不是同一个;但同时由于梯度累积,第一轮完毕后没有执行 ZeroGrad(),第一轮算出来的梯度还存在原 bucket view 里。

综合起来状态就是 grad 非空,grad 的指针和 bucket view 的指针不是同一个,因此会进入背景所述第二条的逻辑里,grad 被指向新的 bucket view,并且标记了下一次直接复写(此时原 bucket view 丢失);但是又由于梯度累积,下一轮还需要接着上一轮的 bucket view 的梯度继续累加,而上一轮的 bucket view 还没来得及用就丢了。

因此造成了反向过程中梯度计算的错误。

修改方式

正确的核心应该是:只要有旧 grad 就 copy 到新 bucket view,否则 mark overwrite。
这个旧 grad 既可能是上述“梯度累积 + rebuild buckets”情况导致的,也可能是用户非要手动修改 tensor->grad 导致的(但就算是这个情况,也应该遵从用户意志,哪怕数学上是不对的,也应该按照用户自己定义的行为执行)。

新的“每一轮 tensor->grad() 重新绑定到新的 bucket view” 的逻辑变为:

  • 如果 grad 为空(ZeroGrad(set_to_none=True) 的结果)
    • 则标记下次覆盖写(因为原先的 bucket view 还有上一轮的梯度)
  • 如果 grad 的指针和 bucket view 的指针不是同一个(RebuildBuckets 的结果
    • 则直接将旧 grad 中的内容 copy 到新的 bucket view 里

@Chamberlain0w0

Copy link
Copy Markdown
Contributor Author
image image

@Chamberlain0w0

Copy link
Copy Markdown
Contributor Author

ctest:
image

@Chamberlain0w0

Copy link
Copy Markdown
Contributor Author

梯度累积修复验证:以测例 3 为基础,测例 3 的配置为八卡 DDP 的 --batch_size 10 --total_batch_size 5120,改成四卡 DDP,但是 bs/gbs 都不变,原则上会进行两步的梯度累积,每步 print 的 loss 仍应该不变。同时保证没开 distopt,ddp 仍走传统 reducer 路径。

八卡 DDP baseline:
image

修复前四卡 DDP(master 分支,错误结果):
image

修复后四卡 DDP(此分支,正确结果):
image

@kilinchange kilinchange merged commit cf21e9f into master Jul 3, 2026
2 checks passed
@kilinchange kilinchange deleted the fix/reducer_grad_accu branch July 3, 2026 08:19
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants