From 37ac206e9ac8972ab802aa66de756cf5cbc2bc45 Mon Sep 17 00:00:00 2001
From: basil-k-aji-dev <70605804+basil-k-aji-dev@users.noreply.github.com>
Date: Mon, 14 Sep 2026 22:22:00 +0530
Subject: [PATCH] [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.
---
 trainer/train_agent.py |  4 +++-
 trainer/train_grpo.py  |  4 +++-
 trainer/train_ppo.py   | 21 +++++++++++++++++++--
 3 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/trainer/train_agent.py b/trainer/train_agent.py
index 50520aa..c8b7f88 100644
--- a/trainer/train_agent.py
+++ b/trainer/train_agent.py
@@ -473,7 +473,9 @@ if __name__ == "__main__":
         Logger('torch.compile enabled')
         rollout_engine.update_policy(model)
     if dist.is_initialized():
-        model = DistributedDataParallel(model, device_ids=[local_rank])
+        # See train_ppo.py: the RoPE buffers are identical on every rank, so
+        # the per-forward buffer broadcast is pure overhead.
+        model = DistributedDataParallel(model, device_ids=[local_rank], broadcast_buffers=False)
     rollout_engine.update_policy(model)
 
     for epoch in range(start_epoch, args.epochs):
diff --git a/trainer/train_grpo.py b/trainer/train_grpo.py
index 0ba7cf5..9d3a232 100755
--- a/trainer/train_grpo.py
+++ b/trainer/train_grpo.py
@@ -314,7 +314,9 @@ if __name__ == "__main__":
         Logger('torch.compile enabled')
         rollout_engine.update_policy(model)
     if dist.is_initialized():
-        model = DistributedDataParallel(model, device_ids=[local_rank])
+        # See train_ppo.py: the RoPE buffers are identical on every rank, so
+        # the per-forward buffer broadcast is pure overhead.
+        model = DistributedDataParallel(model, device_ids=[local_rank], broadcast_buffers=False)
     rollout_engine.update_policy(model)
     
     # ========== 8. 开始训练 ==========
diff --git a/trainer/train_ppo.py b/trainer/train_ppo.py
index 9023c6b..0c8625e 100644
--- a/trainer/train_ppo.py
+++ b/trainer/train_ppo.py
@@ -34,6 +34,19 @@ 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为输出单一价值的线性层
@@ -429,8 +442,12 @@ 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])
+        # 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.
+        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. 开始训练 ==========