- Propagates MoE load-balancing loss through model outputs - Logs task, auxiliary, and weighted losses across strategies - Computes only explicitly requested callback metrics - Preserves tensor compute_loss API and adds regression tests
