Checkpointing: The DistributionGrid also provides an interface to map experts for checkpoint saving
and loading. Based on this interface, utility functions are provided to query and translate model state
dictionary. The utilities are robust to translate parameter mapping from one model distribution strategy
to other while hiding complexity of expert parameter placement from data scientists.
Loss Functions: ORT MoE provides multiple loss functions for flexible MoE experimentation. The
`load_balancing_loss` is the auxiliary loss defined in the Switch Transformers paper to seek uniform
routing of the batch of samples across the experts. The `sparsity_l1_loss` and `mean_importance_loss`
are introduced in the SpeechMoE paper. The former encourages the sparsity of router activation while
the latter leads to a balanced routing strategy. The `z_loss` belongs to the spherical loss family in multi-
Scaling Speech, Language and Vision Models with Mixture of Experts Technique - Microsoft Community Hub