Fix MoE routing and training gradients

#13

Summary

  • match vLLM routing by always using global correction-biased top-k for Sarvam-105B
  • reject n_group or topk_group configuration with a clear error because grouped routing is not supported by this patch
  • keep e_score_correction_bias as persistent checkpoint state without treating it as a trainable parameter
  • restore differentiable single-rank MoE dispatch during training and fail clearly for unsupported expert-parallel training

Validation

  • NVIDIA B200, fixed seed 20260911, vLLM 0.28.0
  • production-shaped router: 128 experts, top-8, 51 tokens; expert IDs matched exactly and maximum routing-weight absolute error was 5.96e-08
  • miniature two-layer Sarvam checkpoint loaded independently by Transformers and vLLM: identical greedy tokens [127, 200, 155, 155]; maximum absolute error over 1,024 vocabulary log-probabilities was 0.0019813
  • CPU training smoke test confirmed finite gradients for inputs, router weights, and expert parameters
rahular changed pull request status to merged

Sign up or log in to comment