Compare commits

...

3 Commits

Author SHA1 Message Date
jingyaogong 7a9137d2e9 [update] trim verbose ddp comments 2026-09-15 14:53:11 +08:00
jingyaogong 7988bc81a3
Merge pull request #855 from basil-k-aji-dev/perf/ddp-followups
[perf] skip the per-forward RoPE buffer broadcast in the RL trainers
2026-09-15 14:51:37 +08:00
basil-k-aji-dev 37ac206e9a [perf] skip the per-forward RoPE buffer broadcast in the RL trainers
Follow-up to the review on #854.

DDP re-broadcasts every registered buffer from rank 0 on each forward. The only
buffers here are freqs_cos / freqs_sin, which each rank computes deterministically
from config and never mutates from data, so the broadcast moves identical bytes
every step for nothing.

Verified with a two-rank CPU/Gloo run: with broadcast_buffers=False the
gradients stay identical across ranks and so do the buffers.

Also document the critic's tied-embedding assumption. CriticModel.forward never
calls lm_head, which is only safe under DDP because tie_word_embeddings (default
True) makes lm_head.weight the same Parameter as embed_tokens.weight. Untied, it
becomes a genuinely unused parameter and DDP raises.
2026-09-14 22:22:00 +05:30
3 changed files with 8 additions and 5 deletions

View File

@ -473,7 +473,8 @@ if __name__ == "__main__":
Logger('torch.compile enabled')
rollout_engine.update_policy(model)
if dist.is_initialized():
model = DistributedDataParallel(model, device_ids=[local_rank])
# 同 train_ppo:RoPE buffer 各 rank 一致,每步广播纯属浪费
model = DistributedDataParallel(model, device_ids=[local_rank], broadcast_buffers=False)
rollout_engine.update_policy(model)
for epoch in range(start_epoch, args.epochs):

View File

@ -314,7 +314,8 @@ if __name__ == "__main__":
Logger('torch.compile enabled')
rollout_engine.update_policy(model)
if dist.is_initialized():
model = DistributedDataParallel(model, device_ids=[local_rank])
# 同 train_ppo:RoPE buffer 各 rank 一致,每步广播纯属浪费
model = DistributedDataParallel(model, device_ids=[local_rank], broadcast_buffers=False)
rollout_engine.update_policy(model)
# ========== 8. 开始训练 ==========

View File

@ -36,7 +36,7 @@ def rep_penalty(text, n=3, cap=0.5):
class CriticModel(MiniMindForCausalLM):
def __init__(self, params):
super().__init__(params)
# 替换lm_head为输出单一价值的线性层
# lm_head 不参与 forward,仅靠 tie_word_embeddings 与 embed_tokens 共享权重,解绑后 DDP 会报未使用参数
self.value_head = nn.Linear(params.hidden_size, 1)
def forward(self, input_ids=None, attention_mask=None, **kwargs):
@ -429,8 +429,9 @@ if __name__ == "__main__":
Logger('torch.compile enabled')
rollout_engine.update_policy(actor_model)
if dist.is_initialized():
actor_model = DistributedDataParallel(actor_model, device_ids=[local_rank])
critic_model = DistributedDataParallel(critic_model, device_ids=[local_rank])
# freqs_cos/freqs_sin 各 rank 由 config 确定性算出,默认每步广播一次纯属浪费
actor_model = DistributedDataParallel(actor_model, device_ids=[local_rank], broadcast_buffers=False)
critic_model = DistributedDataParallel(critic_model, device_ids=[local_rank], broadcast_buffers=False)
rollout_engine.update_policy(actor_model)
# ========== 8. 开始训练 ==========