Implement Router Z-loss

[josejg]

What does this PR do?

Z-loss is an additional term that improves the stability of softmax logit inputs by penalizing large values where the float precision is reduced.

Idea was introduced by the ST-MoE Paper - https://arxiv.org/abs/2202.08906. And has been recently used in the OLMoE work - https://arxiv.org/abs/2409.02060.

What issue(s) does this change relate to?

Alternative implementation of #133

Before submitting

• Have you read the contributor guidelines?
• Is this change a documentation change or typo fix? If so, skip the rest of this checklist.
• Was this change discussed/approved in a GitHub issue first? It is much more likely to be merged if so.
• Did you update any related docs and document your change?
  • Unsure where to document things
• Did you update any related tests and add any new tests related to your change? (see testing)
  • Added a test_moe_forward_backward_with_zloss and test_moe_forward_backward_with_zloss
• Did you run the tests locally to make sure they pass?
• Did you run pre-commit on your change? (see the pre-commit section of prerequisites)

[mihir-db]

LGTM!