Skip to content

[DO NOT REVIEW] Add MXFP8 recipes for the qwen3 and gpt_oss MoE models - #4372

Draft
anijain2305 wants to merge 1 commit into
mxfp8-uneven-shard-paddingfrom
mxfp8-moe-model-configs
Draft

[DO NOT REVIEW] Add MXFP8 recipes for the qwen3 and gpt_oss MoE models#4372
anijain2305 wants to merge 1 commit into
mxfp8-uneven-shard-paddingfrom
mxfp8-moe-model-configs

Conversation

@anijain2305

Copy link
Copy Markdown
Contributor

The grouped-expert support landed without a way to turn it on outside
DeepSeek-V3, so the only model exercising it was the one it was developed
against. Add debug recipes for the two other MoE models whose experts the
converter already handles.

Both quantize the attention projections and the expert grouped GEMMs, and
leave the router gate and lm_head in BF16: the gate is tiny and its output
drives routing, and the vocab projection is large and numerically sensitive.
Only the fused QKV saves a quantized input activation for WGRAD, since
nothing else consumes it, while attention already retains the output
projection's input.

gpt_oss exercises a path DeepSeek-V3 does not: its experts carry per-expert
biases and a SwiGLU clamp, and the converter preserves both by subclassing
GptOssGroupedExperts rather than replacing it -- only the grouped GEMM seam
changes, and the 2D biases are skipped when installing compute weights.

Add CPU coverage for both recipes. MXFP8 needs SM100, which CI does not have,
but the config-tree transform does not, so the test bypasses the hardware
gate and asserts which modules are converted, which input-activation format
each linear saves, and that the experts subclass the model's own variant.

Verified on 2xGB200: qwen3_moe_debug_mxfp8 trains with loss tracking BF16
(6.673 vs 6.603 at step 3) and 0.6 GiB lower peak memory.


Stack created with GitHub Stacks CLIGive Feedback 💬

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Aug 28, 2026
@anijain2305
anijain2305 marked this pull request as draft August 28, 2026 22:44
@anijain2305 anijain2305 changed the title Add MXFP8 recipes for the qwen3 and gpt_oss MoE models [DO NOT REVIEW] Add MXFP8 recipes for the qwen3 and gpt_oss MoE models Aug 28, 2026
@anijain2305
anijain2305 force-pushed the mxfp8-moe-model-configs branch from 0d7ac75 to 18027d9 Compare August 31, 2026 06:00
@anijain2305
anijain2305 force-pushed the mxfp8-moe-model-configs branch from 18027d9 to 89ff3b3 Compare August 31, 2026 06:08
@anijain2305
anijain2305 force-pushed the mxfp8-moe-model-configs branch from 89ff3b3 to fc9b778 Compare August 31, 2026 18:24
@anijain2305
anijain2305 force-pushed the mxfp8-moe-model-configs branch from fc9b778 to f61ae82 Compare August 31, 2026 18:38
@anijain2305
anijain2305 force-pushed the mxfp8-moe-model-configs branch from f61ae82 to e9b4aec Compare August 31, 2026 21:17
@anijain2305
anijain2305 force-pushed the mxfp8-moe-model-configs branch from e9b4aec to d471895 Compare August 31, 2026 21:44
@anijain2305
anijain2305 force-pushed the mxfp8-moe-model-configs branch from d471895 to 0f31ed6 Compare August 31, 2026 22:15
@anijain2305
anijain2305 force-pushed the mxfp8-moe-model-configs branch from 0f31ed6 to 28da47c Compare September 1, 2026 05:43
@anijain2305
anijain2305 force-pushed the mxfp8-moe-model-configs branch from 28da47c to 68ea709 Compare September 1, 2026 05:53
The grouped-expert support landed without a way to turn it on outside
DeepSeek-V3, so the only model exercising it was the one it was developed
against. Add debug recipes for the two other MoE models whose experts the
converter already handles.

Both quantize the attention projections and the expert grouped GEMMs, and
leave the router gate and lm_head in BF16: the gate is tiny and its output
drives routing, and the vocab projection is large and numerically sensitive.
Only the fused QKV saves a quantized input activation for WGRAD, since
nothing else consumes it, while attention already retains the output
projection's input.

gpt_oss exercises a path DeepSeek-V3 does not: its experts carry per-expert
biases and a SwiGLU clamp, and the converter preserves both by subclassing
GptOssGroupedExperts rather than replacing it -- only the grouped GEMM seam
changes, and the 2D biases are skipped when installing compute weights.

Add CPU coverage for both recipes. MXFP8 needs SM100, which CI does not have,
but the config-tree transform does not, so the test bypasses the hardware
gate and asserts which modules are converted, which input-activation format
each linear saves, and that the experts subclass the model's own variant.

Verified on 2xGB200: qwen3_moe_debug_mxfp8 trains with loss tracking BF16
(6.673 vs 6.603 at step 3) and 0.6 GiB lower peak memory.
@anijain2305
anijain2305 force-pushed the mxfp8-moe-model-configs branch from 68ea709 to 7ec2141 Compare September 1, 2026 06:14
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant