From 9c088dd3a57f56de87e70fac02806734f80eae7b 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 14:54:07 +0530 Subject: [PATCH] [fix] route RL trainer forward through DDP so gradients all-reduce train_grpo.py, train_ppo.py and train_agent.py wrap the policy in DistributedDataParallel but run the loss-bearing forward through .module, bypassing the wrapper. DDP arms its gradient all-reduce hooks inside DistributedDataParallel.forward(), so calling .module directly means backward() performs no all-reduce: every rank steps on purely local gradients and the replicas silently diverge. Only rank 0 is saved, so multi-GPU RL training quietly degrades to single-GPU training on 1/N of the data with no error. The supervised trainers are unaffected; they already call the wrapped module. no_grad uses of .module (rollout generation, checkpoint saving) are left untouched, where unwrapping is correct. Single-GPU behaviour is unchanged: with no process group, model is the bare module and these calls are identical to before. --- trainer/train_agent.py | 5 +++-- trainer/train_grpo.py | 5 +++-- trainer/train_ppo.py | 10 +++++----- 3 files changed, 11 insertions(+), 9 deletions(-) diff --git a/trainer/train_agent.py b/trainer/train_agent.py index 261b00c..6b4360f 100644 --- a/trainer/train_agent.py +++ b/trainer/train_agent.py @@ -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) - model_unwrapped = model.module if isinstance(model, DistributedDataParallel) else model 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) logits = res.logits[:, :-1, :] per_token_logps = F.log_softmax(logits, dim=-1).gather(2, input_ids[:, 1:].unsqueeze(-1)).squeeze(-1) diff --git a/trainer/train_grpo.py b/trainer/train_grpo.py index 76ea045..0ba7cf5 100755 --- a/trainer/train_grpo.py +++ b/trainer/train_grpo.py @@ -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] - model_unwrapped = model.module if isinstance(model, DistributedDataParallel) else model 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) per_token_logps = F.log_softmax(res.logits[:, :-1, :], dim=-1).gather(2, outputs[:, 1:].unsqueeze(-1)).squeeze(-1).gather(1, logp_pos) diff --git a/trainer/train_ppo.py b/trainer/train_ppo.py index 07c7738..9023c6b 100644 --- a/trainer/train_ppo.py +++ b/trainer/train_ppo.py @@ -159,8 +159,6 @@ def ppo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, actor_sched clipfrac_sum = 0.0 aux_loss_sum = 0.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): if stop_ppo: 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): 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]) 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) # 在 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_min={torch.exp(_lrv).min().item():.6f} " 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，防止某卡 break 而其它卡继续导致 DDP 死锁