[mxfp8 moe training] bench and profile mxfp8 a2a fwd and bwd separately

[danielvegamyhre]

Stacked PRs:

• ->[mxfp8 moe training] bench and profile mxfp8 a2a fwd and bwd separately #3203
• [mxfp8 moe training] simplify e8m0 -> fp32 calc #3201
• [mxfp8 moe training] integrate triton quant/dequant kernels into mxfp8 all to all #3197
• [mxfp8 moe training] add triton kernel for mxfp8 dequantization #3195

---

[mxfp8 moe training] bench and profile mxfp8 a2a fwd and bwd separately

Single node benchmark on 4xB200 over NVLink shows perf is flat or slightly slower in this setting.

BF16 forward:

MXFP8 forward:

TL;DR is the transfer of the fp8 data is ~1.72x faster than the bf16 data, but the extra overhead of e8m0 scale transfer, quantization and dequantization all adds up to make the perf flat/slower.

input_shape         num_splits    fwd_bf16_ms    fwd_mxfp8_ms    bwd_bf16_ms    bwd_mxfp8_ms
----------------  ------------  -------------  --------------  -------------  --------------
(1, 8192, 5120)              4       0.269697        0.479122       0.695028        0.993789
(2, 8192, 5120)              4       0.347715        0.468697       0.791324        0.872646
(4, 8192, 5120)              4       0.593996        0.585684       1.28176         1.24674
(8, 8192, 5120)              4       1.53808         1.03233        2.38809         2.3224
(16, 8192, 5120)             4       1.77031         1.8789         4.36899         4.46669

Next steps

• If we avoid dequant and stay in mxfp8 through the token shuffle and grouped gemm we can probably get a net speedup. I discussed with @tianyu-l at PTC and we have some early ideas on how this might work w/ torchtitan without being overly intrusive. In the meantime, we need:
  1. mxfp8 token shuffle kernel
  2. update to_mxfp8_then_scaled_grouped_mm to optionally accept scales (pre-quantized tensors)

[pytorch-bot]

🔗 Helpful Links

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

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