[feat] max vio in sft

[Jackmin801]

This PR adds logic to log the max_vio expert balance stat. This is (tokens_per_expert.max() - tokens_per_expert.mean()) / tokens_per_expert.mean() as detailed in eq 4 of https://arxiv.org/pdf/2408.15664v1

---

GitHub Issue: #790
Linear Issue: Resolves PRIMERL-48

[Jackmin801]

Moe logging

Doesnt break non moe logging

SFT

[mikasenghaas]

ah nice, so we are logging an average over num layers here, right? also, how does this work? does tt keep the tokens routed to each expert as a layer attribute? or is this smth you added:)

[Jackmin801]

ah nice, so we are logging an average over num layers here, right? also, how does this work? does tt keep the tokens routed to each expert as a layer attribute? or is this smth you added:)

Yup. Its average over the layers because thats the multiple of slowdown this would cause in EP.

Yep, theres a buffer that stores the num tokens per expert that is added to at each fwd. Its needed for the loss free load balancing that deepseek does. Actually that reminded me I have to reset it to 0 after logging otherwise it will just keep accumulating.

[Jackmin801]

[mikasenghaas]

cleannn, i really like it