[moe training] update tests and benchmarks for torchtitan moe refactor

[danielvegamyhre]

Summary

pytorch/torchtitan#1534 landed in torchtitan, a refactor which provides a generic MoE implementation that can be configured and used by both Llama4 and DeepSeekV3. We need to refactor our MoE training integration tests accordingly.

Test plan

• pytest test/prototype/moe_training/test_training.py -s

Notes

• I replace logger.info calls with print statements because the logging calls don't show up when running pytest tests, even with the "-s" flag. I don't love this but these log statements are only there temporarily for the development process anyway, until the prototype is more mature. If anyone has any better ideas let me know.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/2718

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit e080d55 with merge base 1526dfe ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[danielvegamyhre]

cc @vkuzo @drisspg @HDCharles if one of you could review that would be great