[mxfp8 moe training] migrate ep utils to mx_formats

[danielvegamyhre]

Prev PR in stack:

• [mxfp8 moe training] move mxfp8 grouped mm code into mx_formats #3898

Summary

• migrate mxfp8 ep utils to mx_formats:
  • from: torchao/prototype/moe_training/ep/
  • to: torchao/prototype/mx_formats/expert_parallel/
• intended final package name after migration + refactor:
  • torchao.training.mxfp8.expert_parallel
• make autograd func imports lazy to avoid circular imports
• move GroupedMMConfig base class to mx_formats

Note/question

• We can call this package something more descriptive than ep, open to suggestions. One idea:
  • expert_parallel (final name after migration would be

Tests

• ./test/prototype/moe_training/test_everything.sh

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 11868c7 with merge base 4a42d32 ():
💚 Looks good so far! There are no failures yet. 💚

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