[mxfp8 moe training] add MX MoE model converter using torchao mxfp8 moe training

[danielvegamyhre]

Context

• We've recently landed prototype mxfp8 MoE training support in torchao, with ~1.3x speedup over bf16 in single node. fsdp2 only, llama4 e2e training with seq_len 8192, 2 experts per device: [mxfp8 moe training] use new 3d colwise quantization kernel ao#3037.
• This PR integrates this feature into torchtitan.

Changes

• Remove MoE conversion logic from mx model converter. Using low precision training on MoE layers should not be dependent on converting dense layers (plus, putting all the logic in one converted was getting messy).
• Add new model converter mx_moe, with a user facing API to convert MoE layers to use mxfp8 scaled grouped gemms, instead of high precision grouped gemms.
• This clean separation allows users more flexibility/control, to use mxfp8 for dense layers, MoE layers, or both.

Test plan

• NGPU=2 CONFIG_FILE="./torchtitan/experiments/llama4/train_configs/debug_model.toml" ./run_train.sh --parallelism.data_parallel_shard_degree=2 --parallelism.tensor_parallel_degree=1 --model.converters="mx_moe" --mx_moe.fqns="experts" --model.print-after-conversion --metrics.log_freq=10 --training.steps=30 --compile.enable

[danielvegamyhre]

cc @tianyu-l @drisspg for review

[drisspg]

it feels like the converter registered for self.config should handle this case specifically

[danielvegamyhre]

as in having separate converters registered for mxfp8 moe, fp8 moe etc? If so, I kind of agree actually, the converters are becoming a bit of a mess, and having separate converters would also allow users to convert just dense or just MoE (or both), rather than the current state of having to convert dense in order to convert MoE.

[drisspg]

Yeah that is one way, another thing I was thinking is that https://github.com/pytorch/ao/blob/93030e750186ace1c1c2ee7a849e2818a9f0ffde/torchao/prototype/moe_training/conversion_utils.py#L50 should be able to gracefully handle the case where the module has already be converted

[danielvegamyhre]

Separated the dense and MoE converters

[tianyu-l]

In the log

Swapped w1.weight to ScaledGroupedMMTensor

it's hard to tell which w1 is converted. Can we include full fqn?

(eager only currently - compile support in progress)

Does it mean we shouldn't care too much about performance with this PR?
How about numerics?

I'll try to find some time to look into #1651

[danielvegamyhre]

@tianyu-l @drisspg this is ready for another look

[danielvegamyhre]

it's hard to tell which w1 is converted. Can we include full fqn?

Yes, i've been meaning to make this update in torchao... will do.

[tianyu-l]

so you prefer to not update the mx_moe configs in the toml configs?

[tianyu-l]

Separating Linear and MoE treatment is reasonable. I have two comments:

1. To make naming consistent, should we rename MX to MXLinear?
2. In general, as we are creating more and more quantization dataclasses in job_config.py, I somehow prefer we group them under the same root level dataclass, say Quantization, matching the folder organization. I think this can be achieved by using hierarchical dataclasses and nested toml https://toml.io/en/v1.0.0#array-of-tables. We may have to update config/manager.py so doesn't need to be addressed in this PR.

[danielvegamyhre]

To make naming consistent, should we rename MX to MXLinear?

That's fine with me, can do in follow up

I somehow prefer we group them under the same root level dataclass, say Quantization, matching the folder organization. I think this can be achieved by using hierarchical dataclasses and nested toml

I like this idea!

[tianyu-l]

How about float8, should we also put the constant here?

[danielvegamyhre]

Yes, I plan to refactor float8 rowwise MoE to use the same pattern in a self-contained follow up PR.

[vkuzo]

I think the torchtitan feature name here should match the torchao one. If you want mxtraining, make the torchao workflow adhere to that too.