MoE code fixes

[gbesposito]

What does this PR do?

This PR fixes some general issues that were in the MoE code

General Changes

• Removes torchtitan dependency by integrating a torch-native Expert Parallelism class, inspired by torchtitan's
• Fixes a bug in the MixedPrecision policy that was keeping experts' parameters' in FP32 instead of casting them
• Replaces the TP-dimension workaround by introducing a first-class EP dimension into the device mesh

Checklist before submitting final PR

• My PR is minimal and addresses one issue in isolation
• I have merged the latest version of the target branch into this feature branch
• I have reviewed my own code w.r.t. correct implementation, missing type hints, proper documentation, etc.
• I have run a sample config for model training
• I have checked that all tests run through (python tests/tests.py)
• I have updated the internal changelog (CHANGELOG_DEV.md)

[rrutmann]

LGTM :)