transformers

7460d407 - Fix NaN gradients in expert-parallel training: mask uninitialized grouped_mm rows

Fix NaN gradients in expert-parallel training: mask uninitialized grouped_mm rows Under EP, sentinel token-expert slots sit beyond offsets[-1] and torch._grouped_mm leaves those output rows (fwd output and bwd d_input) uninitialized. The forward relied on a single post-mask plus a single pre-mask, letting NaN/Inf from uninitialized memory transit the activation and down-projection backward. The gate product's backward (act_fn(gate) * up) turns 0 x Inf into NaN (torch.autograd anomaly mode names this exact Mul), and it escapes into finite gradients: full fine-tuning of any EP-sharded MoE produced nan grad_norm on the second step (the first step survives only because freshly-allocated CUDA memory happens to be zeroed) and the loss collapsed to 0. Mask the sentinel-tail rows after each grouped GEMM instead. Full fine-tuning of OLMoE-1B-7B under ep=4 now matches the single-GPU loss trajectory.