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。

[chen2021673]

这里return nullptr就好了吧

[chen2021673]

PyTorch 在这里和 Megatron 不太一样，进入 no_sync() 时保存旧状态，退出时会恢复旧状态，所以能支持嵌套作用域。目前我觉得不太需要，如果后续需要支持嵌套的话可以参考。

[chen2021673]

这里主要是给 Pipeline 每个 chunk 注册回调函数，但 PipelineSchedule 自己能访问 stage_->chunks，能不能让 PipelineSchedule 直接遍历 chunks 创建 guard 啊，这样相关的逻辑不用放在训练入口。

[kilinchange]

stage 里维护的 chunk 是 Module 基类类型，需要根据 ddp_world_size 判断，才能安全将 module 转换成 DistributedDataParallel 类型。因此不建议目前在 pp 模块内部引入 ddp 相关的耦合逻辑，后续训练入口这块可以考虑统一整理成类似 megatron train.py 的形式，提供一个统一的训练入口。

[kilinchange]

stage 里维护的 chunk 是 Module 基类类型，需要根据 ddp_world_size 判断，才能安全将 module 转换成 DistributedDataParallel 类型。因此不建议目前在 pp 模块内部引入 ddp 相关的耦合逻辑，后续训练入口这块可以考虑统一整理成类似 megatron train.py 的形式，提供一个统一的训练入口。