[moe] brings batch/sequence-wise load balance loss#2061

Open

rakkit wants to merge 5 commits into

pytorch:mainfrom

rakkit:moe_load_balance_loss

Commits on Nov 19, 2025

1) make the moe's load_balance_coeff configurable 2) add the batch and seq-wise aux loss for load balance
rakkit
committed

Commits on Dec 9, 2025

fix loss build and move loss-weights to training scope, and other fixing
rakkit
committed

Commits on Dec 10, 2025