[moe training] Add 2D parallel (FSDP2 + TP) tests for routed experts

[danielvegamyhre]

Stack

• [WIP] [moe training] Add tests for 3D parallel (FDSP + TP + EP) for routed experts #2481
• [moe training] Add 2D parallel (FSDP2 + TP) tests for routed experts #2475 <- this PR
• [moe training] Add TP support for routed experts #2473
• [moe training] Cast to mixed precision policy param dtype in fsdp_pre_all_gather hook #2455

Summary

• New integration test uses NGPU=4, with FDSP=2 and TP=2. Applies FSDP2 first, then applies TP.
• Note all tests currently only apply float8 to routed experts, not shared experts.

Test plan

• ./test/prototype/moe_training/test_fsdp_tp.sh

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

⏳ 1 Pending, 1 Unrelated Failure

As of commit 80c88fc with merge base 6821971 ():

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Run TorchAO Experimental Tests / test-mps-ops (macos-m1-stable) (gh) (trunk failure)
  Process completed with exit code 127.

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

[vkuzo]

at this point I'd also add test_everything.sh which runs all of the shell scripts, same as we do for float8 training

[danielvegamyhre]

MPS test failure is unrelated to this change