[QUESTION] How to enable fp8 dispatch while training MOE models?

[new-TonyWang]

Your question
How can I enable FP8 dispatch for MoE model training? I've found the FP8 Linear and GroupLinear implementations, but I don't see any code covering FP8 All-to-All communication. DeepEP is capable for dispatch FP8 tensor, but it seems megatron does not pass fp8 tensor to DeepEP.
megatron code
deepep interface code
deepep require pass a tuple(fp8_tensor, fp8_scale) to dispatch fp8 tensor, but megatron only pass one tensor

[yaox12]

It's not supported yet and there's no plan to support it. We evaluated FP8 dispatch some times before, but we couldn't get any perf gain from it. The reason is that FP8 dispatch only sends the row-wise quantized tokens, then you need another "de-quantize and re-quantize along the other dimension" kernel to get the column-wise quantized tensors for backward computation. As a result, the elapsed time of (FP8 dispatch + de-quantize-and-re-quantize) is almost the same as or even longer than BF16 dispatch.