[roadmap/tracker] Low precision MoE training

[danielvegamyhre]

Creating this issue as a roadmap/tracker for enabling float8 training for MoEs with token-choice routing. Both core requirements as well as ideas for additional performance optimizations are included.

Compute

• fp8 rowwise
  • Add torch._scaled_grouped_mm kernel in core
    • bf16 grouped gemm pytorch#150374 (by @ngimel)
  • Add differentiable scaled grouped mm with dynamic float8 rowwise quant in torchao
    • Initial prototype of differentiable _scaled_grouped_mm function #1969
  • Add custom kernels in torchao for performing per-group scaling on device, to avoid host-device sync
    • [scaled grouped mm] add triton kernels for float8 rowwise quantization with per-group/jagged scales #2064
    • [scaled grouped mm] integrate triton kernels into differentiable scaled grouped mm #2077
  • Fuse padding of group sizes up to nearest multiple of 16 into the dynamic quant kernel
• mxpf8
  • mxfp8 scaled grouped gemm Add MXFP8 Support to scaled_grouped_gemm pytorch#153502
  • torchao differentiable _scaled_grouped_mm support for mxpf8 recipe for dynamic quant before grouped GEMMs
  • triton kernels to do scaling per group without d2h sync

Communication

I looked at traces and validated "all to all dispatch -> grouped gemm -> all to all combine" are all sequentially dependent, so in theory faster/low precision comms should improve performance. There is some overlap with the shared expert computation, but it is not 100% overlap, so there is room for optimization. This will be especially important if/when "all to all" spans multiple nodes, where inter-node network bandwidth is lower than the intra-node NVLink bandwidth.

This is also inspired by the DeepSeekV3 paper where, if I understand correctly, they do a2a dispatch in fp8 but keep a2a combine in bf16 as they found it was more sensitive to low precision during training.

• Add on device all_to_all_v kernels compatible with:
  • mxfp8 (P0)
  • float8 rowwise (P1)
• token permutation kernel supports low precision dtypes by permuting scales to be in proper order for permuted tokens (link)
  • mxfp8 (P0)
  • float8 rowwise (P1)

Torchao UX

• Add tensor subclass (ScaledGroupedMMTensor) with an op override for torch.aten._grouped_mm => runs differentiable scaled grouped mm
  • float8 moe training conversion API prototype #2275
• Add one line model conversion API, should recursively swap nn.Parameter data tensors of the expert weights with ScaledGroupedMMTensor.
  • float8 moe training conversion API prototype #2275
• support configurable recipe (fp8 rowwise, mxpf8)

Compile support

• Compile support for torch._grouped_mm
  • done by @bdhirsh in basic compile support for grouped_mm pytorch#153384
• Differentiable _scaled_grouped_mm can compile with fullgraph=True
• E2E compilation of each TranformerBlock in torchtitan after MoE conversion via tensor subclass approach

Distributed support

• Composability with FSDP2 (will likely need something like this for the new tensor subclass)
  • mxfp8 (P0)
  • float8 rowwise (P1) [float8 moe training] FSDP support  #2413
• Composability with TP
  • mxfp8 (P0)
  • float8 rowwise [float8 moe training] Add TP support #2425
• Composability with FSDP + TP 2D parallel
  • mxfp8 (P0)
  • float8 rowwise (P1) [float8 moe training] Add TP support #2425
• Composability with tp2ep
  • mxfp8 (P0)
  • float8 rowwise (P1)
• Composability with dp2ep
  • mxfp8 (P0)
  • float8 rowwise (P1)

[danielvegamyhre]

cc @tianyu-l @ngimel @drisspg if you have any feedback on the direction/ideas proposed here

[yiakwy-xpu-ml-framework-team]

Hi @danielvegamyhre have your team added quantization tools to convert deepseek bf16 (sft version) to fp8 with block_wise scalars? (better to be aligned with deepseek quantization format Dyanmic FP8 quant and config.json, then I can use it in SGLang and vLLM directly)

I try to seek help to llmcompressor, but they are not ready. Wish updates from torch trainning team if this feature has been supported.

[danielvegamyhre]

Hi @danielvegamyhre have your team added quantization tools to convert deepseek bf16 (sft version) to fp8 with block_wise scalars?

@yiakwy-xpu-ml-framework-team we are currently planning to focus on mxfp8 next due to high interest from various users we've talked to, but fp8 blockwise is on our radar as well! We have an inference only deepseek style blockwise fp8 prototype. If you want to work on a PR to support training, we'd happily work with you on reviewing and merging that. Otherwise, we'll work on it as soon as we can :)

cc @vkuzo @lessw2020