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.
execute_tool() relied on signal.SIGALRM, which does not exist on Windows.
The AttributeError was swallowed by a bare except, so every tool call
silently returned None and the model always observed
{"error": "tool not found"} (issue #846).
calculate_math also ran eval() on model-generated text. Two of the three
call sites (scripts/eval_toolcall.py, scripts/web_demo.py) passed no
globals at all, leaving full builtins reachable from sampled output.
safe_math_eval() in trainer_utils.py parses the expression with ast and
walks it, allowing only arithmetic operators and a math whitelist, with
caps on expression length and exponent size. It cannot loop forever, so
the signal-based timeout is no longer needed and is removed. Both sqrt(x)
and math.sqrt(x) spellings keep working.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Same fused AdamW optimization as the full SFT script, using the
existing device_type variable so the CPU-only path is unaffected.
Co-authored-by: TRAE CLI <traecli@bytedance.com>
Use PyTorch's fused AdamW implementation to reduce optimizer kernel
launch overhead and speed up the training step on CUDA devices.
Guard with device check so the CPU-only fallback path is unaffected.
Co-authored-by: TRAE CLI <traecli@bytedance.com>
Move log_softmax for mb_resp_logp inside autocast_ctx so it is upcast
to fp32, matching the rollout path (compute_per_token_logps runs under
no_grad+autocast). Previously log_softmax ran outside autocast on
fp16/bf16 logits, introducing a ~1e-2 systematic bias between
mb_resp_logp and old_resp_logp, so ratio != 1 and KL != 0 even before
any weight update. Under sparse rule-based rewards (e.g. math PPO)
this bias destabilises training.
Also add:
- --rollout_eval_mode (default 1): switch actor to eval during rollout
to disable dropout and keep old/new logp consistent.
- --debug_log_ratio: optional diagnostic printing |mb-old| on the first
minibatch to verify ratio ~= 1.
Neil Ge
jingyaogong
jingyaogong
basil-k-aji-dev
jingyaogong
tkgcc
neil.ge
Justin Guese
FloatFrank
TKiteRunner
jingyaogong