[mxfp8 moe traing] to_mxfp8_a2a_dequant flat perf vs high precision a2a

[danielvegamyhre]

Looking at torchtitan traces, it seems the mxfp8 a2a dispatch is ~2x faster than bf16 (1681us vs 3487us), but the mxfp8 a2a combine is (1) roughly exact same duration as bf16 a2a combine, and (2) bf16/mxfp8 a2a impls both take roughly 23x longer than a2a dispatch. This is unexpected.

Tlparse: https://manifold.edge.x2p.facebook.net/v0/read/tree/logs/.tmpHhVPih/rank_0/index.html?bucketName=tlparse_reports&apiKey=tlparse_reports-key&withPayload=1&timeoutMsec=10000

Trace for llama4 EP=4, using defalut a2a_impl

Trace for llama4 EP=4, to_mxfp8_a2a_dequant

[danielvegamyhre]

Looking at the backward pass, I see 6 all2alls, in groups of 3 each, each with 2 very small ones and 1 very large one.

In the bf16 run, for each group, the 1st short a2a is not overlapped with the FDSP reduce-scatter of the grads. This results in them both being very fast:

In the mxfp8 a2a run, for each group, the 1st short a2a IS overlapped with the FDSP reduce-scatter of the grads. This results in them both being very slow:

[danielvegamyhre]

I guess one thing I need to do is a set a manual seed to try to make token routing consistent between the bf16 a2a run and mxfp8 a2a run.

[danielvegamyhre]

Update: i thought this could be due to the straggler affect, where there is expert load imbalance across devices, so perhaps a large part of the 80ms seen in the trace is just waiting for other devices to arrive at the same point.

To test this, i hardcoded equal expert load balance in torchtitan and verified it's working, but I still don't see the speedup i expect

[danielvegamyhre]

Today, letting training run for 200+ steps seems to produce better results than 50 steps. I get 5-7% speedup with mxfp8 a2a, although I still would expect more, so more investigation to do here. See numbers in this PR: pytorch/torchtitan#1765