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:96426bd9303825371ba6798545cbed28f4abe961
Branches Tags
github:master
github:website
github:v2
github:minimind-v1
..
compare: github:937e4174a4b3896c42fdb901d8a43181506c66b5
Branches Tags
github:master
github:website
github:v2
github:minimind-v1

4 Commits
96426bd930 ... 937e4174a4

Author SHA1 Message Date
Yixian Chen 937e4174a4
Merge 138c367ddb into 7a9137d2e9 2026-09-16 09:21:13 +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

@ -473,7 +473,8 @@ if __name__ == "__main__":
Logger('torch.compile enabled') Logger('torch.compile enabled')
rollout_engine.update_policy(model) rollout_engine.update_policy(model)
if dist.is_initialized(): 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) rollout_engine.update_policy(model)
for epoch in range(start_epoch, args.epochs): for epoch in range(start_epoch, args.epochs):

3
trainer/train_grpo.py
View File

@ -314,7 +314,8 @@ if __name__ == "__main__":
Logger('torch.compile enabled') Logger('torch.compile enabled')
rollout_engine.update_policy(model) rollout_engine.update_policy(model)
if dist.is_initialized(): 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) rollout_engine.update_policy(model)
# ========== 8. 开始训练 ========== # ========== 8. 开始训练 ==========

7
trainer/train_ppo.py
View File

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