diff --git a/trainer/train_agent.py b/trainer/train_agent.py index c8b7f88..63a0782 100644 --- a/trainer/train_agent.py +++ b/trainer/train_agent.py @@ -473,8 +473,7 @@ if __name__ == "__main__": Logger('torch.compile enabled') rollout_engine.update_policy(model) if dist.is_initialized(): - # See train_ppo.py: the RoPE buffers are identical on every rank, so - # the per-forward buffer broadcast is pure overhead. + # 同 train_ppo:RoPE buffer 各 rank 一致,每步广播纯属浪费 model = DistributedDataParallel(model, device_ids=[local_rank], broadcast_buffers=False) rollout_engine.update_policy(model) diff --git a/trainer/train_grpo.py b/trainer/train_grpo.py index 9d3a232..ad9aa3b 100755 --- a/trainer/train_grpo.py +++ b/trainer/train_grpo.py @@ -314,8 +314,7 @@ if __name__ == "__main__": Logger('torch.compile enabled') rollout_engine.update_policy(model) if dist.is_initialized(): - # See train_ppo.py: the RoPE buffers are identical on every rank, so - # the per-forward buffer broadcast is pure overhead. + # 同 train_ppo:RoPE buffer 各 rank 一致,每步广播纯属浪费 model = DistributedDataParallel(model, device_ids=[local_rank], broadcast_buffers=False) rollout_engine.update_policy(model) diff --git a/trainer/train_ppo.py b/trainer/train_ppo.py index 0c8625e..4d762b6 100644 --- a/trainer/train_ppo.py +++ b/trainer/train_ppo.py @@ -34,22 +34,9 @@ def rep_penalty(text, n=3, cap=0.5): # 自定义的Critic模型,继承自MiniMindLM class CriticModel(MiniMindForCausalLM): - """Value head over the base model. `forward` never calls `lm_head`. - - This relies on `tie_word_embeddings` (default True), which makes - `lm_head.weight` the *same* Parameter as `model.embed_tokens.weight`. The - embedding is used on every forward, so the tied parameter still receives a - gradient and DDP sees no unused parameters. - - Untie it and `lm_head.weight` becomes a real parameter that this forward - never touches, so DDP raises "Expected to have finished reduction in the - prior iteration" unless the critic is wrapped with - `find_unused_parameters=True` (or `lm_head` is deleted from the critic). - """ - 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): @@ -442,10 +429,7 @@ if __name__ == "__main__": Logger('torch.compile enabled') rollout_engine.update_policy(actor_model) if dist.is_initialized(): - # broadcast_buffers=False: the only registered buffers are the RoPE - # tables (freqs_cos / freqs_sin), which every rank recomputes - # deterministically from config and never mutates from data. The - # default re-broadcasts them from rank 0 on every forward for nothing. + # 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)