exclude top k from max vio calculation to try and skip the "padding experts"

[Jackmin801]

Note

Updates get_load_balance_stats to optionally ignore top-k routed experts (avoiding padding experts) when computing max violation, with proper buffer reset.

• MoE Load Balance Stats (src/prime_rl/trainer/model.py):
  • Add try_to_avoid_padding_experts: bool = True to get_load_balance_stats.
  • When enabled, sort mlp.tokens_per_expert desc and slice off router.top_k to exclude top-k routed experts from mean/max calc.
  • Ensure reset zeros the original mlp.tokens_per_expert buffer.
  • Add explicit tensor type annotation for tokens_per_expert.

^{Written by Cursor Bugbot for commit a1d15ea. This will update automatically on new commits. Configure here.}

[mikasenghaas]

this would make us blind to very heavy expert imbalance where all tokens are routed to a singular expert?

[Jackmin801]

yea it would but also its better than us not caring about the metric because we always think its just padding tokens ruining the metric