Fix fp16 underflow in MoE load balancing loss by enforcing fp32 softmax

[ITcarrot]

What does this PR do?

This PR updates the load_balancing_loss_func across various MoE models to explicitly perform the softmax operation in float32.

When running models in fp16, aux_loss underflow to 0 when using large batchsize. The underflow comes from doing the flowing operation in fp16:

# Compute the average probability of routing to these experts
router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum(
        router_per_expert_attention_mask, dim=0
)

Casting fp32 in softmax can ensure a precise and robust load balancing loss calculation.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@Cyrilvallez @IlyasMoutawwakil

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: dbrx, ernie4_5_moe, ernie4_5_vl_moe, flex_olmo, glm4v_moe, gpt_oss, granitemoe, granitemoehybrid, granitemoeshared, jamba, jetmoe, minimax, minimax_m2, mixtral, olmoe, phimoe