[moe training] Add TP support for routed experts

[danielvegamyhre]

Stack

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

Summary

• Adds TP integration test and testing utils file for shared functions used in single GPU / FSDP / TP training tests
• Update _scaled_grouped_mm to support 3D A tensor, which is needed for the shared expert (code).
• Make offs optional, to handle shared_expert case where num_experts=1 (no group offsets needed since there's only 1 token group).

Test plan

• ./test/prototype/moe_training/test_tp.sh

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure, 1 Pending

As of commit b9e58fe with merge base 01f7352 ():

NEW FAILURE - The following job has failed:

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

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

[vkuzo]

is the bf16 group gemm working on A100? If yes, I would vote for adding an emulation mode and running this test in emulation mode. We do this for float8 and MX training.

[danielvegamyhre]

I will look into this and add emulation if bf16 grouped gemm builds on a100

[vkuzo]

nit: why do we need recursion to check this?

[danielvegamyhre]

It's just a generic way of checking all target FQNs were converted properly, and verifying all non-target FQNs were correctly not converted. It can easily be applied when we extend tests to other MoE models as well beyond just the torchtitan llama4 one I started with.