Repository navigation
fix: fix param grad sync under SP - #224
Conversation
| } else { | ||
| optimizer = optimizer_creator(named_parameters); | ||
| } | ||
| optimizer->set_model_grad_finalizer(nn::parallel::FinalizeModelGrads); |
There was a problem hiding this comment.
FinalizeModelGrads 绑定到 optimizer 上感觉职责不太合适。optimizer 接收到的 grad 应该已经是 finalize 后可用于更新的完整梯度,本身只负责后续的 clip / step 等操作。
另外,将 FinalizeModelGrads 放到 optimizer->Step() 内部也容易影响 clip 执行顺序(理论上应该保证:
FinalizeModelGrads -> clip -> step,因为 clip 需要基于最终同步完成的梯度计算)。
建议目前在 optimizer->Step() 之前显式调用 nn::parallel::FinalizeModelGrads(named_parameters);后续如果抽象统一的 training/schedule 入口,再将这部分逻辑收进去。Megatron 也是在 forward/backward schedule 结束后调用 finalize_model_grads,而不是绑定到 optimizer:
(no pipeline schedule 情况,带 pipeline 时类似)https://github.com/NVIDIA/Megatron-LM/blob/0ac6ffd3859fef41bbfcd92a67d56861a79d4b34/megatron/core/pipeline_parallel/schedules.py#L866
There was a problem hiding this comment.
改了一下。现在是 training loop 里面需要显式调用 nn::parallel::FinalizeModelGrads,但是传入的是 model chunks vector。在 backward 完成、optimizer 更新前,统一调用 FinalizeModelGrads(model_chunks):先完成 DP 同步,再执行 SP 梯度的 TP SUM。
另外对应做的修改:
-
原位于 DistOpt 里面的 StartGradSync() 和 FinishGradSync() 交由 DDP model 管理;同时删除 StartGradSync(),因为 DDP 里面实际同步开启时机相对固定,简化起见目前不需要额外留这个用户接口;
-
未做引入 main_grad 基建的设计,而是让 bucket group 新增 LocalGradShard 映射表:DistOpt 在原有 BuildShardParamsAndBindGrads() 中登记原参数与本地梯度 view,finalizer 据此处理 ZeRO-1/2,尤其解决 ZeRO-2 原参数 grad() 为空时的梯度访问问题。只保存引用,不新增显存 buffer。
问题简述:
我们现在实现没有一个统一的 main_grad 的全局 buffer,并且 zero2 里面切梯度,只在反向过程中会算临时 full grad,一方面不会把这个 full grad 写入 tensor.grad,RS 之后也不会把 grad shard 写到 tensor.grad 里面(不符合 tensor.shape == grad.shape 的约束)。之后 finalize 统一入口没法通过 param 获取对应的 param.grad,可能需要额外的维护一个 map
e8f6bf2 to
df9871b
Compare
df9871b to
7ed52c9
Compare
7ed52c9 to
c15c2d1
Compare








https://gxtctab8no8.feishu.cn/docx/FR88dJOZPoXT1Ox1TH1c2abpnof