transformers

4a0dc8fb - Fix wrong gradients for all non-expert parameters in expert-parallel training

Fix wrong gradients for all non-expert parameters in expert-parallel training Under EP the router hook zeroes the routing scores of non-local experts, so in backward each rank's score gradient covers only the slots of its local experts, and nothing sums the per-rank partial gradients: the gate weights and, through the gate's input, every parameter upstream of each MoE block receive gradients missing the contributions that flow through remote experts. The existing _AllReduceBackward on the experts' hidden input covers the dispatch branch, and the top_k_weights branch is explicitly skipped when is_expert_parallel -- but under EP it is exactly as partial as under TP-MoE. Measured against a single-GPU reference (OLMoE-1B-7B, one batch, fp32 so rounding noise vanishes): before the fix, 3/179 parameters agree (relative max-abs errors 0.3-2.5 on attention, norms, embeddings and router gates, 10-100x above the run-to-run noise floor; only the last layer's experts and the final norm -- the parameters backward reaches before crossing an expert block -- are correct). After the fix: 179/179 agree, max relative error 2.7e-5. Fix: allreduce-sum the score gradient in the EP router hook, before the non-local mask (each slot has exactly one owning rank, so the sum is exact).