MoE training: is there an API/hook to monitor expert activations?

[tiesanguaixia]

Hi Slime team, thanks for open-sourcing this project.

I’m training an MoE model with Slime and would like to monitor expert activations during training (e.g., per-expert token counts / routing distribution / load balance stats). Does Slime provide any built-in interface, hook, or recommended way to collect and log these metrics (e.g., to W&B)?

Any pointers to docs or the right place in the codebase to add instrumentation would be appreciated!