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

Merge pull request #854 from basil-k-aji-dev/fix/rl-ddp-grad-sync 

[fix] route RL trainer forward through DDP so gradients all-reduce
pull/780/merge
jingyaogong 2026-09-14 22:28:03 +08:00 committed by GitHub
parent a3c7b01cc0 9c088dd3a5
commit c367572702
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 11 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
trainer/train_agent.py
View File

@ -278,9 +278,10 @@ def rl_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_model
rewards = calculate_rewards(prompts, completions, gt_batch, tools_batch, args.num_generations, reward_model, device=args.device, turn_outputs_batch=turn_outputs_batch, unfinished_batch=unfinished_batch) rewards = calculate_rewards(prompts, completions, gt_batch, tools_batch, args.num_generations, reward_model, device=args.device, turn_outputs_batch=turn_outputs_batch, unfinished_batch=unfinished_batch)
model_unwrapped = model.module if isinstance(model, DistributedDataParallel) else model
with autocast_ctx: with autocast_ctx:
res = model_unwrapped(input_ids, attention_mask=full_mask) # 反向传播必须经过 DDP 包装后的模块：直接调用 .module 会跳过
# DDP 的 prepare_for_backward，梯度不会 all-reduce，各卡静默发散。
res = model(input_ids, attention_mask=full_mask)
aux_loss = res.aux_loss if lm_config.use_moe else torch.tensor(0.0, device=args.device) aux_loss = res.aux_loss if lm_config.use_moe else torch.tensor(0.0, device=args.device)
logits = res.logits[:, :-1, :] logits = res.logits[:, :-1, :]
per_token_logps = F.log_softmax(logits, dim=-1).gather(2, input_ids[:, 1:].unsqueeze(-1)).squeeze(-1) per_token_logps = F.log_softmax(logits, dim=-1).gather(2, input_ids[:, 1:].unsqueeze(-1)).squeeze(-1)

5
trainer/train_grpo.py
View File

@ -95,9 +95,10 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod
rewards = calculate_rewards(prompts, completions, reward_model).to(args.device) # [B*num_gen] rewards = calculate_rewards(prompts, completions, reward_model).to(args.device) # [B*num_gen]
model_unwrapped = model.module if isinstance(model, DistributedDataParallel) else model
with autocast_ctx: with autocast_ctx:
res = model_unwrapped(outputs, attention_mask=full_mask) # 反向传播必须经过 DDP 包装后的模块：直接调用 .module 会跳过
# DDP 的 prepare_for_backward，梯度不会 all-reduce，各卡静默发散。
res = model(outputs, attention_mask=full_mask)
aux_loss = res.aux_loss if lm_config.use_moe else torch.tensor(0.0, device=args.device) aux_loss = res.aux_loss if lm_config.use_moe else torch.tensor(0.0, device=args.device)
per_token_logps = F.log_softmax(res.logits[:, :-1, :], dim=-1).gather(2, outputs[:, 1:].unsqueeze(-1)).squeeze(-1).gather(1, logp_pos) per_token_logps = F.log_softmax(res.logits[:, :-1, :], dim=-1).gather(2, outputs[:, 1:].unsqueeze(-1)).squeeze(-1).gather(1, logp_pos)

10
trainer/train_ppo.py
View File

@ -159,8 +159,6 @@ def ppo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, actor_sched
clipfrac_sum = 0.0 clipfrac_sum = 0.0
aux_loss_sum = 0.0 aux_loss_sum = 0.0
log_count = 0 log_count = 0
actor_unwrapped = actor_model.module if isinstance(actor_model, DistributedDataParallel) else actor_model
critic_unwrapped = critic_model.module if isinstance(critic_model, DistributedDataParallel) else critic_model
for ppo_epoch in range(args.ppo_update_iters): for ppo_epoch in range(args.ppo_update_iters):
if stop_ppo: if stop_ppo:
break break
@ -168,11 +166,13 @@ def ppo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, actor_sched
for i in range(0, B, mb_size): for i in range(0, B, mb_size):
inds = b_inds[i:i + mb_size] inds = b_inds[i:i + mb_size]
mb_values_seq = critic_unwrapped(input_ids=gen_out[inds], attention_mask=full_mask[inds]) # 反向传播必须经过 DDP 包装后的模块：直接调用 .module 会跳过
# DDP 的 prepare_for_backward，梯度不会 all-reduce，各卡静默发散。
mb_values_seq = critic_model(input_ids=gen_out[inds], attention_mask=full_mask[inds])
mb_resp_values = mb_values_seq.gather(1, logp_pos[inds]) mb_resp_values = mb_values_seq.gather(1, logp_pos[inds])
with autocast_ctx: with autocast_ctx:
res = actor_unwrapped(input_ids=gen_out[inds], attention_mask=full_mask[inds]) res = actor_model(input_ids=gen_out[inds], attention_mask=full_mask[inds])
aux_loss = res.aux_loss if lm_config.use_moe else torch.tensor(0.0, device=args.device) aux_loss = res.aux_loss if lm_config.use_moe else torch.tensor(0.0, device=args.device)
# 在 autocast 内计算 log_softmax，避免直接对 fp16/bf16 logits # 在 autocast 内计算 log_softmax，避免直接对 fp16/bf16 logits
# 计算造成额外数值偏差。 # 计算造成额外数值偏差。
@ -191,7 +191,7 @@ def ppo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, actor_sched
f"ratio_max={torch.exp(_lrv).max().item():.6f} " f"ratio_max={torch.exp(_lrv).max().item():.6f} "
f"ratio_min={torch.exp(_lrv).min().item():.6f} " f"ratio_min={torch.exp(_lrv).min().item():.6f} "
f"dropout={getattr(lm_config, 'dropout', None)} " f"dropout={getattr(lm_config, 'dropout', None)} "
f"training={actor_unwrapped.training}") f"training={actor_model.training}")
approx_kl = (0.5 * (log_ratio ** 2) * resp_policy_mask[inds]).sum() / resp_policy_mask[inds].sum().clamp(min=1) approx_kl = (0.5 * (log_ratio ** 2) * resp_policy_mask[inds]).sum() / resp_policy_mask[inds].sum().clamp(min=1)
# 同步各卡的 approx_kl，防止某卡 break 而其它卡继续导致 DDP 死锁 # 同步各卡的 approx_kl，防止某卡 break 而其它卡继续导致 DDP 死锁