[WIP] [moe training] Add tests for 3D parallel (FDSP + TP + EP) for routed experts

[danielvegamyhre]

WIP because it requires pytorch/pytorch#157216 to land first

Stack

• [WIP] [moe training] Add tests for 3D parallel (FDSP + TP + EP) for routed experts #2481 <- this PR
• [moe training] Add 2D parallel (FSDP2 + TP) tests for routed experts #2475
• [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

• Test applies 3D parallelism to routed experts, where FSDP=2, TP=2, EP=2

Test plan

• ./test/prototype/moe_training/test_fsdp_tp_ep.sh

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ You can merge normally! (1 Unrelated Failure)

As of commit cfef1fb with merge base 2defe30 ():

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.

[danielvegamyhre]

cc @tianyu-l I would appreciate your feedback on the parallelism configuration for this test.

Note:

1. This is all using your PR for local testing: dp2ep Expert Parallel torchtitan#1324
2. The target model is just importing the Llama4 MoE layer from torchtitan directly, not the full transformer.
3. apply_moe_ep_tp is a modified version of the torchtitan function, which just removes the for-loop over transformer blocks (since there is no transformer), but is otherwise identical.
4. I'm attempting to test something like a single-host dp2ep configuration.

Is this correct/useful or do you recommend any adjustments?