[float8 moe training] FSDP support

[danielvegamyhre]

Summary

• Remove use_triton_for_per_group_scales flag, as it was making the code too complicated/ugly for not enough value - I can bench torch vs triton impls by manually changing 2 lines, no need for this config flag.
• Update the ScaledGroupedMMTensor subclass to support FSDP ops while preserving the subclass.
• Since the subclass wraps the nn.Parameter data tensor, it's subclass must be preserved through FSDP related ops (otherwise, if the ops produce regular tensors, the grouped_mm => differentiable scaled grouped mm op override won't happen.

Design

• Keep "grouped_mm => differentiable scaled grouped mm" override in torch_function to capture graph for autograd
• For all other ops, handle in torch_dispatch to avoid autograd
• For ops related to FSDP, preserve subclass. Otherwise, returning a regular tensor is fine.

Test plan

Integration tests

• Added FSDP integration test in test/prototype/moe_training/test_fsdp.py

Manual tests

• Tested with llama4 on 2 H100s w/ FSDP=2, using torchtitan PoC integration, added temporary logging to confirm the op override differentiable scaled grouped mm is indeed being called in the MoE.
• Command: CUDA_VISIBLE_DEVICES="6,7" 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" 2>&1 | tee ~/ao/fsdp-log.txt

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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! (2 Unrelated Failures)

As of commit 57d600e with merge base 809af2e ():

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

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

• Run Regression Tests / test (CPU 2.7, linux.4xlarge, torch==2.7.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh) (trunk failure)
  test/quantization/pt2e/test_numeric_debugger.py::TestNumericDebuggerInfra::test_prepare_for_propagation_comparison
• Run Regression Tests / test (CUDA 2.7, linux.g5.12xlarge.nvidia.gpu, torch==2.7.0, cuda, 12.6) / linux-job (gh) (trunk failure)
  test/quantization/pt2e/test_numeric_debugger.py::TestNumericDebuggerInfra::test_prepare_for_propagation_comparison

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

[danielvegamyhre]

fyi @tianyu-l @ngimel as well for awareness on progress