fix: fix param grad sync under SP - #224
Chamberlain0w0 wants to merge 1 commit into
Conversation
| DataType Dtype() const; | ||
|
|
||
| void set_sequence_parallel(bool enabled) { sequence_parallel_ = enabled; } | ||
| bool sequence_parallel() const { return sequence_parallel_; } |
There was a problem hiding this comment.
这里得根据 parallel::global::GetSequenceParallelEnabled() 决定是否 set 吧
| } 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
2 participants
https://gxtctab8no8.feishu.cn/docx/FR88dJOZPoXT1Ox1TH1c2abpnof