[moe training] update tests + benchmarks with conditional runs based on SM arch; make test cases more comprehensive and consistent

[danielvegamyhre]

Summary

• Add conditional logic to benchmarks and tests to only run if the SM arch supports the underlying GEMMs that are going to run (when applicable), to address float8 rowwise scaled grouped mm doesn't support B200 #2904 (torch._scaled_grouped_mm not supported on B200 yet)
• Make all MoE training tests have identical comprehensive parameterization where applicable, for robustness + consistency. All tests now have cases for all combinations of:
  • recipes: mxfp8, fp8 rowwise
  • execution: eager, compiled
  • target fqns: routed experts only, routed experts + shared experts
  • distributed tests cover FDSP, TP, FSDP+TP. EP tests coming next.
• Add TFLOPs and speedup calculations to grouped gemms bench script
• In benchmark_scaled_grouped_mm_dq.py split into (1) fwd time, and (2) e2e fwd+bwd time. Backward 2d-2d GEMM is not ready yet so I want to measure + optimize forward.
• Add more tests to test_everything.sh and ensure everything passes.

Test plan

• ./test/prototype/moe_training/test_everything.sh

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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:

• ROCm CI/CD workflows failing due to : download from https://api.github.com/repos/pytorch/pytorch timed out.

⏳ No Failures, 7 Pending

As of commit 7eb6bfd with merge base 083d0c3 ():
💚 Looks good so far! There are no failures yet. 💚

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

[vkuzo]

stamping prototype code