[MXFP8 MoE] What's the expected inference solution on H100s, after training with TorchAO MXFP8 MoE?

[goldhuang]

Hi team,

Thanks for your great implementation of the new MXFP8 MoE! I have integrated it and consider to use it for prod training.
But I got a concern about how to do inference.

MXFP8 is only available on B200. What is the expected inference solution on H100 or even non-Nvidia GPUs after training with MXFP8. Other quantizations, even another FP8 quantization, is not guaranteed to work well with the model trained with MXFP8.

Is a QAT finetuning with another quantization method expected?
Should we just inference with another quantization method without finetuning?

I guess FP4 training is a similar case.

I think the question is not only to TorchAO team. Anyone please share your ideas/insights if you would like to.

Thanks in advance!

[danielvegamyhre]

@goldhuang how you are doing inference currently (vllm, sglang, etc)?

[goldhuang]

@danielvegamyhre We do standalone inference deployment, without vllm or sglang.
For example, do model expert by AOTInductor then run the package during inference. I think that's the regular inference deployment for most of non-LLM models.

[danielvegamyhre]

@goldhuang my initial thoughts are:

• You probably don't need QAT unless you want to inference with a totally different quantization method, like fp8 blockwise or int8.
• Using mxfp8 dynamic activation + dynamic weight quantization would be easiest to get good results (evals) with, since this is the same exact quantization method used for training, so weights are optimized to be robust against the kind of slight information loss in this quantization method. This will provided limited speedup though, compared to other methods.
• Prequantizing the weights to mxfp8 then use weight-only quantization is probably a reasonable strategy to start with, and should provide better speedups than ^.

[goldhuang]

@danielvegamyhre mxfp8 is only available on B200. My initial question is the expected inference solution on H100 or even non-Nvidia GPUs after training with mxfp8.
Other quantizations, even another FP8 quantization, is not guaranteed to work well with the model trained with MXFP8.

[danielvegamyhre]

@goldhuang i see, then yes, for serving on different hardware with a different quantization strategy i would recommend doing QAT for target method for inference (e.g. float8 per-row weight-only), to fine-tune the weights to be robust to the inference quantization you'll use.

QAT docs:

• README with simple example: https://github.com/pytorch/ao/blob/main/torchao/quantization/qat/README.md
• Detailed walkthrough: https://docs.pytorch.org/ao/main/finetuning.html
• APIs: https://docs.pytorch.org/ao/main/api_ref_qat.html

Note QAT apis in torchao only support nn.Linears now, we don't support 3d nn.Parameters yet (which the MoE training code is designed for). So the overall MoE support in torchao is still in progress for some stages of the model lifecycle

[jcaip]

FYI The current state of MoE inference in torchao:

• 3d params / non-weight params are okay after landing Add quantize_ nn.Parameter support #3083, so you should be able to quantize non-linear MoE models that use for-loop linear for forwards
• bmm support for Float8Dynamic rowwise is currently being worked on Enable PerRow(axis) to support axes other than -1 #3303 but once this land you can quantize models that use torch.bmm for moe forwards.
• for models that used grouped_mm, I don't think this is supported by our APIs currently