[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.

UPDATE 07/22/2025: revised priorities to reflect shifting focus to prioritize mxfp8

This is not an exhaustive list, but highlights some primary milestones / requirements

Compute

• mxpf8
  • mxfp8 scaled grouped gemm
    • 2d-3d gemm for output and dX ([mxfp8 moe] add support for fbgemm 2d-3d mx8mx8bf16 grouped gemm #2848)
    • 2d-2d gemm for dW (Add 2d-2d support to MXFP8 Grouped GEMM FBGEMM#4816)
  • torchao differentiable _scaled_grouped_mm support for mxpf8 recipe for dynamic quant before grouped GEMMs
  • triton kernels for per token group scale conversion to blocked swizzled format
    • for 2d inputs with groups along M ([mxfp8 moe training] add per group blocked scale kernels #2886)
    • for 3d expert weights ([mxfp8 moe training] add triton kernel for blocked swizzled 3d weight scales #2894)
    • for 2d in puts with groups along K ([mxfp8 moe training] per group scale conversion to blocked format with groups along K dim (for 2d2d grouped gemm) #2956)
• 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
  • Faster inductor codegen kernels for dynamic quant of 3d tensors along dim1: Inductor codegen for float8 dynamic quantization ops for scaled_grouped_mm backward pass is slow pytorch#159769
    • alternatively, handwritten triton kernel faster than torch.compile for this ([moe training] add fp8 rowwise kernels for expert weights #2696)
      • this also needs to be faster [MoE fp8 rowwise training] Runtime of quantizing 3d expert weights scales worse than linearly #2880
• fp8 blockwise
  • quant primitives
  • Explore DeepGEMM for fp8 blockwise grouped GEMM
  • triton kernels to do scaling per group without d2h sync

Communication

I looked at traces and validated "all to all dispatch and shuffle -> grouped gemm -> all to all combine and unshuffle" 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.

• mxfp8
  • 1d on device all_to_all_v comms kernel (differentiable; with dynamic quant) [mxfp8 moe training] mxfp8_on_device_all_to_all_v kernel #3048
  • NVSHMEM put/get primitives instead of tl.load/tl.store to support inter-node a2a
• float8 blockwise (P0)

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 blockwise/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 (fullgraph=False)
• E2E compilation of each TranformerBlock in torchtitan after MoE conversion via tensor subclass approach (fullgraph=True)

Distributed support

• Composability with FSDP2 (will likely need something like this for the new tensor subclass)
  • mxfp8 (P0)
  • float8 blockwise (P0)
  • float8 rowwise (P1) [float8 moe training] FSDP support  #2413
• Composability with TP
  • mxfp8 (P0)
  • float8 blockwise (P0)
  • float8 rowwise (P1) [moe training] Add TP support for routed experts #2473
• Composability with FSDP + TP
  • mxfp8 (P0)
  • float8 blockwise (P0)
  • float8 rowwise (P1) [moe training] Add 2D parallel (FSDP2 + TP) tests for routed experts #2475
• Composability with dp2ep as implemented here: dp2ep Expert Parallel torchtitan#1324
  • mxfp8 (P0)
  • float8 blockwise (P0)
  • float8 rowwise (P1) [WIP] [moe training] Add tests for 3D parallel (FDSP + TP + EP) for routed experts #2481

[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

[danielvegamyhre]

Added pytorch/pytorch#159769 to track inductor codegen improvements we'll need for fp8 rowwise backward. Hoping after mxfp8 and fp8 blockwise we'll have time to get fp8 rowwise in good shape as well for PTC.