[float8 moe training] fix bug affecting mixed precision training

[danielvegamyhre]

Fixes pytorch/torchtitan#1332

Summary

Previously when testing the MoE training prototype manually with Llama4 in torchtitan, I had to add model = model.to(torch.bfloat16) in torchtitan's train.py to workaround an error resulting from what I thought was a bug in torchtitan not applying bf16 mixed precision policy for FDSP2 to MoE layers properly:

    File "/home/danvm/torchtitan/torchtitan/experiments/llama4/model/moe.py", line 83, in forward
      x.dtype == self.w1.dtype == self.w2.dtype == self.w3.dtype == torch.bfloat16
  AssertionError: torch._grouped_mm only supports bf16 dtypes

However, after debugging further I determined the bug was actually in the torchao ScaledGroupedMMTensor subclass not setting the dtype properly, the way WeightWithDynamicFloat8CastTensor does, where for ops that preserve the subclass, dispatch returns a ScaledGroupedMMTensor using dtype pulled while unwrapping the args/kwargs.

Reasoning: as a general rule for pytorch, the dtype of the output will be the same as one of its inputs (with the exception of a handful of ops with the 'out_dtype' arg). So we are enforcing a simple version of this to ensure weights are cast to the mixed precision dtype properly in FSDP.

Other

• Add some logging
• Add helper script for running FSDP tests which require torchrun w/ some args

Test plan

• torchao tests: ./test/prototype/moe_training/test_fsdp.sh
• torchtitan manual testing: NGPU=2 CONFIG_FILE="./torchtitan/experiments/llama4/train_configs/debug_model.toml" ./run_train.sh --training.steps=100 --model.converters="float8" --float8.moe_fqns_prototype="experts"

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• VolumeLimitExceeded Issue for linux.2xlarge and linux.4xlarge

⏳ No Failures, 6 Pending

As of commit 4675cda with merge base 994a4ba ():
💚 Looks good so far! There are no failures yet. 💚

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

[danielvegamyhre]

cc @drisspg @vkuzo for review