[mxfp8 moe training] temp workaround: don't compile GroupedExperts

[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

---

[mxfp8 moe training] temp workaround: don't compile GroupedExperts

See thread (#2250 (comment)) for context.

TL;DR is this avoids a tensor metadata mismatch issue between forward output and backward() input (upstream grad).

Tests

CONFIG_FILE=/home/dev/torchtitan/torchtitan/models/deepseek_v3/train_configs/deepseek_v3_16b.toml ./run_train.sh --metrics.log_freq=10 \
--training.steps=1500  \
--parallelism.data_parallel_shard_degree=4 \
--parallelism.expert_parallel_degree=4 \
--parallelism.tensor_parallel_degree=2 \
--parallelism.expert_tensor_parallel_degree=1 \
--training.seq_len=8192 \
--activation_checkpoint.mode=full \
--model.print_after_conversion \
--training.local_batch_size=16 \
--quantize.linear.mx.mxfp8_dim0_cast_kernel_choice="triton" --quantize.linear.mx.mxfp8_dim1_cast_kernel_choice="cuda" \
--quantize.grouped_mm.mx.fqns="experts" --quantize.grouped_mm.mx.recipe_name="mxfp8_wgrad_with_hp" \
--compile.enable --compile.components="model,loss" --debug.moe_force_load_balance \
--model.converters="quantize.grouped_mm.mx"

[bdhirsh]

there's also another hardcoded call to compile _run_experts_grouped_mm here that you'd need to skip (I'm doing it in my patch in https://gist.github.com/bdhirsh/970a671b84c35cc95a76f33657ca4d69)

[tianyu-l]

revert

[tianyu-l]

can we remove this logic since we are not compiling GroupedExperts any more?

[tianyu-l]

add a comment on the issue this code is working around

[danielvegamyhre]

@tianyu-l thanks for the early feedback. the top 3 PRs in this stack aren't quite ready - once scale testing is complete and i've confirmed the mxfp8 a2a feature works as expected (perf, numerics) then i'll polish it up, add PR descriptions, and let you know when it's ready.

So far so good, just waiting for a long training run to finish to verify training stability and identical convergence to bf16. Should be ready early next week.

[tianyu-l]

we don't have to land this, if #2281 is going to be landed soon