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.
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.
jingyaogong
jingyaogong
basil-k-aji-dev
TOPSAIL\78190
TKiteRunner
jingyaogong
readlnh