Normalising a single top-k weight by its own sum gives exactly 1.0 for every
token, so the gate leaves the graph entirely: topk_idx is a discrete index and
the weight is now a constant, which leaves the LM loss no differentiable path
to the router. It then trains on aux_loss alone, whose optimum is a uniform
distribution, and the experts never specialise.
Measured on MOEFeedForward with router_aux_loss_coef=0, so only the LM path can
contribute a gradient:
num_experts_per_tok=1, norm_topk_prob=True 7.02e-10 <- current default
num_experts_per_tok=1, norm_topk_prob=False 1.02e-02
num_experts_per_tok=2, norm_topk_prob=True 1.67e-02
Only normalise when there is more than one weight to normalise.
Note this changes the forward pass for existing num_experts_per_tok=1 MoE
checkpoints: the expert output was scaled by a constant 1.0 and is now scaled
by the gate probability.
Closes#858.
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>
Neil Ge
jingyaogong
jingyaogong
Daoyuan Li
Linxiushen
basil-k-aji-dev
wbbeyourself
jingyaogong
tkgcc
neil.ge
Justin Guese
TOPSAIL\78190