fix: fix grad accumulation/overlap under ZeRO-2

[Chamberlain0w0]

背景

先前的实现 is_last_microbatch 的处理过于草率，导致目前 ZeRO-2 同时开启梯度累积和 overlap_grad_reduce 时，当前实现会在第一个 microbatch backward 期间发起 reduce-scatter，并将 grad_reduce_dispatched_ 设置为 true。

该状态直到 optimizer->step() 调用 FinishGradSync() 后才会重置，因此存在两个问题：

1. 后续 microbatch 不会再次发起梯度规约，新增梯度无法参与本轮 reduce-scatter。
2. 第一个 microbatch 的异步 reduce-scatter 可能尚未结束，后续 microbatch 已开始写入 temp_full_grad_buffer，造成计算与通信之间的数据竞争。

关闭 overlap_grad_reduce 时不会触发该问题，因为梯度同步统一在所有 microbatch 完成后的 optimizer->step() 中执行。

修改内容

参考 Megatron-LM 的 no_sync 机制，引入真实的 is_last_microbatch_ 控制：

• 为 Module 增加通用的 no_sync() 接口和 RAII NoSyncGuard。
• DistributedDataParallel::no_sync() 在 guard 生命周期内将 bucket group 的 is_last_microbatch_ 设置为 false，退出时恢复为 true。
• 非最后一次 backward 仅将梯度累积到本地 buffer，不标记 bucket ready，也不发起 reduce-scatter。
• 最后一次 backward 恢复梯度同步，使完整的累积梯度只进行一次规约。
• 未开启 overlap_grad_reduce 时保持原有行为，由 optimizer->step() 发起同步。
• Pipeline Parallel 通过注入的 no_sync_func_ 使用该机制，不直接依赖或包含 DDP 实现。
• PP/VPP 根据每个本地 model chunk 的实际 backward 顺序确定最后一次 backward，避免将 microbatch ID 错误地等同于 backward 完成顺序。
• GPT-2 和 LLaMA3 的非 PP 训练循环在非最后一个梯度累积 microbatch 中持有 NoSyncGuard。
• 不新增 overlap_grad_reduce 命令行参数，继续使用 DDP 配置中的默认行为。

行为变化

开启梯度累积和 overlap_grad_reduce 后，同一个 optimizer step 内的执行过程变为：

1. 非最后一个 microbatch：仅累积本地梯度。
2. 最后一个 microbatch：bucket ready 后发起一次异步 reduce-scatter。
3. optimizer->step()：等待通信完成并更新参数。

这样可以确保 reduce-scatter 读取的是所有 microbatch 累积后的完整梯度，同时避免通信期间继续修改 full gradient buffer。

[kilinchange]

实现放到 .cc 里。

[Chamberlain0w0]

改了

[kilinchange]

实现放到 .cc 里。

[Chamberlain0w0]

改了

[kilinchange]

我看注释里说如果只在 last microbatch sync 的话，这里就不需要 clear 了，但下面的 clear 没有删，确认下是否需要删除？

[Chamberlain0w0]

可以删，这块是忘删了。

[kilinchange]

实现放到 .cc 里

[Chamberlain0w0]

改了

[kilinchange]

前置声明 NoSyncGuard 就行，这里不需要 include module.h

[Chamberlain0w0]

改了

[kilinchange]

zero_stage=0 且开启梯度累积时，仍然会所有 microbatch 都进行梯度通信，但实际上也只需要最后一个 microbatch 发起通信，看看这个 pr 里方不方便一起改了？不好改的话就留个 FIXME 吧

[Chamberlain0w0]

统一解决了一下，现在三条 DDP 的路径都完成了在梯度累积下的 last microbatch 再发起梯度同步：

1. naive DDP（不带分桶）：给 DDP class 里加了一个 atomic_bool 的 is_last_microbatch_ 成员，用来标记是否为 last mb；DDP::SetIsLastMicrobatch() 会更新它；然后修改 AllReducePostAccumulateHook， 让其额外接受一个 bool 的参数，来决定是否要真的执行 allreduce。

AllReducePostAccumulateHook(
    function::ReduceOpType reduce_op,
    const ProcessGroup *pg = nullptr,
    std::shared_ptr<const std::atomic_bool> enabled = nullptr);

2. 带分桶 DDP：给 Reducer class 里加了一个 bool 的 is_last_microbatch_ 成员，用来标记是否为 last mb；DDP::SetIsLastMicrobatch() 会更新它；BucketHook 入口 Reducer::MarkVariableReadyDense 处判断此 bool，来决定是否要继续执行标记 grad ready 开启通信的逻辑。

3. DistOpt-based ZeRO DDP（本 PR 原本核心改动）：给 BucketGroup class 里面加了一个 bool 的 is_last_microbatch_ 成员，用来标记是否为 last mb；DDP::SetIsLastMicrobatch() 会更新它；准备更新梯度的时候，在ParamAndGradBucketGroup::RegisterGradReady 入口进行判断，决定是否要进行梯度同步。

三条路径运行时互斥，实际运行时只会存在一种 DDP 的实现方式，所以 DDP::SetIsLastMicrobatch() 并不会产生冲突；在存在梯度累积时，三条 DDP 路径都通过本 PR 中核心修改的同一个 NoSyncGuard 来控制“非最后 microbatch 不通信”。