github
/
minimind
mirror of https://github.com/jingyaogong/minimind.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Compare commits

base: github:0f1eb4a4ae5d5ba46302fdc59679def285dd9459
Branches Tags
github:master
github:website
github:v2
github:minimind-v1
..
compare: github:8acbed85ff3828f5c8aa9cd0d53f7aeb6707e5ec
Branches Tags
github:master
github:website
github:v2
github:minimind-v1

4 Commits
0f1eb4a4ae ... 8acbed85ff

Author SHA1 Message Date
AmirX0 8acbed85ff
Merge 5b26153b75 into 7a9137d2e9 2026-09-15 19:34:49 +08:00
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
Show Stats Expand all files Collapse all files

3
trainer/train_agent.py
View File

@ -469,7 +469,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):

3
trainer/train_grpo.py
View File

@ -309,7 +309,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. 开始训练 ==========

7
trainer/train_ppo.py
View File

@ -31,7 +31,7 @@ warnings.filterwarnings('ignore')
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):
@ -424,8 +424,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. 开始训练 ==========