Add MXFP8 fused-MLP overrides in torchtitan

[wolfcomos]

Summary

Adds an opt-in override module, torchtitan/overrides/mxfp8_fused_mlp.py, that runs a FeedForward / routed-experts MLP entirely in MXFP8 with a fused SwiGLU boundary: one composite autograd function computes w13 GEMM -> fused SwiGLU + dual MXFP8 quantization -> w2 GEMM (dense and grouped variants), using torchao's fused gated-activation kernel (pytorch/ao#4743) for the activation boundary and torchao's MXFP8 grouped/scaled-mm internals for the GEMMs. Compared to the unfused MXFP8 path, the fused boundary eliminates the standalone activation-quantization casts (1 fwd + 1 bwd fused op instead of 6 + 6 standalone cast kernels per MLP) and never materializes the BF16 post-activation tensor.

The integration:

• MXFP8FusedMLP and MXFP8FusedGroupedMLP subclass the stock FeedForward / GroupedExperts and keep stock parameters (w1/w2/w3); the two projection weights are packed into the kernel's fused layout at forward time. No parameter surgery, no state-dict hooks, no init/sharding remaps — checkpoints and fresh initialization are bitwise-identical to stock modules (asserted by a shipped test).
• Two @override factories: mxfp8_fused_mlp (targets FeedForward.Config) and mxfp8_fused_grouped_mlp (targets RoutedExperts.Config, which owns both the token dispatcher and the inner experts — the factory swaps the dispatcher to the padded variant its kernels require, pad_multiple=128).
• Fail-loud, no silent fallback: non-stock or already-converted configs, DTensor dense activations, biased projections, and missing kernels raise with actionable messages.
• Zero changes to existing files' behavior: overrides/fused_swiglu.py, components/quantization/mx.py, and the quantization converters are untouched, and this module imports nothing from the BF16 fused-swiglu override. Enabling the override is the only opt-in; additional MXFP8 fusion paths (e.g. fully fused grouped MLPs) extend this module with their own composite and factory.

Usage

# dense (llama3 debugmodel example flavor)
--override.imports torchtitan.overrides.mxfp8_fused_mlp.mxfp8_fused_mlp

# routed experts (deepseek_v3 debugmodel example flavor; needs EP>=2 so the token
# dispatcher produces the padded expert-major layout the kernels consume)
--override.imports torchtitan.overrides.mxfp8_fused_mlp.mxfp8_fused_grouped_mlp

Example flavors: llama3_debugmodel_mxfp8_fused_mlp, deepseek_v3_debugmodel_mxfp8_fused_mlp (attention/lm_head stay BF16 — composable with MXFP8LinearConverter on those modules independently).

Numerics and tests

• Shipped: tests/unit_tests/test_mxfp8_fused_mlp.py — 8 CPU-runnable tests (config-tree transform, apply_overrides, meta-device build, dispatcher pad_multiple, factory fail-loud, checkpoint keys == stock, and a seeded fresh-init state-dict bitwise comparison against the stock modules), validated with GPUs visible, with CUDA_VISIBLE_DEVICES="", and on stock PyPI torchao without a GPU (all 8 execute).
• Because parameters and checkpoint keys are stock, checkpoints interchange freely with unfused modules in both directions.
• GPU numerics are validated in NVIDIA-internal Blackwell CI and maintained there (upstream CI has no SM100): 25-test suite on GB200 covering bitwise forward parity between the fused kernel and standalone-cast quantization, SQNR tracking against the per-GEMM MXFP8 and BF16 references, torch.compile with unbacked routing-dependent token counts, dispatcher pad-row/tail inertness, and profiler op-count contracts.
• Kernel-level numerics (bitwise forward, one-code-bounded backward at ~6e-7 of elements, root-caused to the kernel's correctly-rounded sigmoid + FMA contraction) are documented in [mxfp8 training] Add a fused gated-activation (SwiGLU) + MXFP8 quantization kernel ao#4743.

🤖 Generated with Claude Code