[mxfp8 moe training] support wgrad_with_hp recipe

[danielvegamyhre]

Stacked PRs:

• [mxfp8 moe training] temp workaround: don't compile GroupedExperts #2268
• [mxfp8 training] add new configurable params now exposed by torchao #2251
• [mxfp8 moe training] mxfp8 all to all #2250
• ->[mxfp8 moe training] support wgrad_with_hp recipe #2249

---

Summary

torchao MXFP8 MoE training code now supports a new recipe: wgrad_with_hp (described below). This PR update the torchtitan integration to allow users to use. it

Recipes:
- "mxfp8": Use MXFP8 for all 3 grouped GEMMs in the forward and backward pass (output, dgrad, wgrad).
- "mxfp8_wgrad_with_hp": Use MXFP8 for forward output and dgrad, but keep wgrad in high-precision.
This can be used to trade-off some performance for improved accuracy.

Note: I plan to do some benchmarking to provide more concrete guidance to users on what expert shapes will result in better TPS using wgrad_with_hp

Tests

mxfp8 recipe:

CONFIG_FILE=/home/dev/torchtitan/torchtitan/models/llama4/train_configs/llama4_17bx16e.toml ./run_train.sh --metrics.log_freq=10 \--training.steps=200  \--parallelism.data_parallel_shard_degree=8 \--parallelism.expert_parallel_degree=8 \--parallelism.tensor_parallel_degree=1 \                                                  
--parallelism.expert_tensor_parallel_degree=1 \
--profiling.enable_profiling --profiling.profile_freq=30 \
--training.seq_len=8192 \
--activation_checkpoint.mode=none \
--model.print_after_conversion \
--training.local_batch_size=12 \
--model.converters="quantize.grouped_mm.mx,quantize.linear.mx" \
--quantize.linear.mx.mxfp8_dim1_cast_kernel_choice="cuda" \
--quantize.linear.mx.filter_fqns="output,router.gate,wk,wv" \
--quantize.grouped_mm.mx.fqns="experts" --quantize.grouped_mm.mx.recipe_name="mxfp8_wgrad" \
--compile.enable --debug.moe_force_load_balance

mxfp8_wgrad_with_hp recipe:

CONFIG_FILE=/home/dev/torchtitan/torchtitan/models/llama4/train_configs/llama4_17bx16e.toml ./run_train.sh --metrics.log_freq=10 \--training.steps=200  \--parallelism.data_parallel_shard_degree=8 \--parallelism.expert_parallel_degree=8 \--parallelism.tensor_parallel_degree=1 \                                                  
--parallelism.expert_tensor_parallel_degree=1 \
--profiling.enable_profiling --profiling.profile_freq=30 \
--training.seq_len=8192 \
--activation_checkpoint.mode=none \
--model.print_after_conversion \
--training.local_batch_size=12 \
--model.converters="quantize.grouped_mm.mx,quantize.linear.mx" \
--quantize.linear.mx.mxfp8_dim1_cast_kernel_choice="cuda" \
--quantize.linear.mx.filter_fqns="output,router.gate,wk,wv" \
--quantize.grouped_mm.mx.fqns="experts" --quantize.grouped_mm.mx.recipe_name="mxfp8_wgrad_with_hp" \
--compile.enable --debug.moe_force_load_balance

[danielvegamyhre]

cc @tianyu-l for review. pre-commit passes locally but fails in CI, not sure why yet. I did a fresh pip install of requirements-dev.txt

[tianyu-l]

Would love to see some requests / tech reports / experiments / scientific studies / guidance on such new features -- with all these flexible options, it's not clear to users what works best.

[danielvegamyhre]

Would love to see some requests / tech reports / experiments / scientific studies / guidance on such new features -- with all these flexible options, it's not clear to users what works best.

We are publishing docs and an e2e tutorial on torchao docsite on this, which we can link to in the torchtitan MXFP8 docs, if that works? If so, I can add a new PR on top of this stack with documentation updates.

[tianyu-l]

We are publishing docs and an e2e tutorial on torchao docsite on this

would love to read this tutorial

[danielvegamyhre]

We are publishing docs and an e2e tutorial on torchao docsite on this

would love to read this tutorial

Here is the tutorial preview: https://docs-preview.pytorch.org/pytorch/ao/3752/mxfp8_expert_parallel.html

the torchtitan config params will need to be updated since i changed the MXFP8ExpertParallel to be applied automatically for mxfp8_wgrad_with_hp recipe, but otherwise it is good to go i think, let me know if you have any thoughts

[danielvegamyhre]

test failures are unrelated

[tianyu-l]

In the tutorial

The mxfp8_wgrad_with_hp recipe is required for MoE training with expert parallelism.

• why it's only required for EP?
• why here default is not wgrad with hp?
• why mxfp8 is an option if hp wgrad is "required"?

[tianyu-l]

I just saw #2250 (comment)

Does it make sense that we don't give user control, just do

• mxfp8 when DeepEP is used, or no EP is used
• mxfp8_wgrad_with_hp when EP > 1 but DeepEP is not used

When EP is not enabled, which one should user use?

[danielvegamyhre]

When EP is not enabled, which one should user use?

• For performance, it depends on the expert shapes, batch size, and sequence length. For directional guidance on which recipe is better in a given context, I'm thinking about working on some tables like we have for float8 (see screenshot below). There has been positive feedback from users on this.
• For accuracy / improved step quality, wgrad_with_hp compute weight gradients in bf16, so this recipe can be more of a net benefit, not in terms of TPS but in terms of "time to target validation loss" or "time to some eval score threshold." Luca also found perf benefit for certain smaller shapes, with fp8_rowwise_with_gw_hp (Hopper recipe). These have to evaluated through experimentation though.