[mxfp8 moe training] update benchmarks to force load balancing

[danielvegamyhre]

Force expert load balancing to simulate steady state MoE training (after early steps load balancing stabilizes).

New benchamarks:

Llama4 16e:

CUDA_VISIBLE_DEVICES=7 python benchmarks/prototype/moe_training/bench_moe_layer.py --recipe mxfp8 --local_num_
experts 1 --seq_len 8192 --local_batch_size 16 --hidden_dim 8192 --dim 5120
total_M: 131072, N: 8192, K: 5120
bf16 time: 252.799 ms
mxfp8 time: 185.525 ms
speedup: 1.363x

DSV3 671b:

CUDA_VISIBLE_DEVICES=7 python benchmarks/prototype/moe_training/bench_moe_layer.py --recipe mxfp8 --local_num_experts 1 --seq_len 8192 --local_batch_size 16 --hidden_dim 2048 --dim 7168
total_M: 131072, N: 2048, K: 7168
bf16 time: 96.285 ms
mxfp8 time: 83.764 ms
speedup: 1.149x

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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! (1 Unrelated Failure)

As of commit 8740a0f with merge base b644211 ():

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

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

• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh) (trunk failure)
  test/dtypes/test_nf4.py::TestComm::test_comm

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