[moe training] Cast to mixed precision policy param dtype in fsdp_pre_all_gather hook

[danielvegamyhre]

Stack

• [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 <- this PR

Summary

• After examining this example and discussing with @weifengpy, we think that when implementing fsdp pre/post all gather hooks for a tensor subclass, the developer must handle casting params to the MP policy param dtype in the pre all-gather hook themselves.
  • I thought the downcasting to save network bandwidth via bf16 all gather instead of fp32 would be done automatically by fsdp, but it appears when implementing these hooks we have to handle it ourselves.
  • Tracking issue to determine contract between MP policy and tensor subclass extension: [FSDP2] figure out the contract for mp_policy and tensor subclass extention pytorch#157395
• Update fsdp_post_all_gather_hook to more correctly handle case where out != None (see code comments for details)
• Remove dtype param from ScaledGroupedMM tensor, as it is no longer needed when doing the casting in pre all gather rather than post all gather.

Test plan

• ./test/prototype/moe_training/test_fsdp.sh
• Manual test with torchtitan llama4 debug model: NGPU=2 CONFIG_FILE="./torchtitan/experiments/llama4/train_configs/debug_model.toml" ./run_train.sh --training.steps=10 --model.converters="float8" --float8.recipe_name="rowwise" --float8.moe_fqns_prototype="experts"

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

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

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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, 5 Pending

As of commit bb9626e with merge base ac14d92 ():

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.

[danielvegamyhre]

cc @drisspg @vkuzo for review

[vkuzo]

do we have a test for this?

[danielvegamyhre]

We have a test for float MoE + FSDP training. We don't have a test verifying which code branch is followed in this fsdp_post_all_gather hook at training step 0 vs 1, but I think the FSDP test alone is sufficient. Let me know if you have other thoughts.

[danielvegamyhre]

MPS test failures are unrelated to this change