[update] trim verbose ddp comments

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)
 

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)
     

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)