[DO NOT REVIEW] Add MXFP8 recipes for the qwen3 and gpt_oss MoE models - #4372
Draft
anijain2305 wants to merge 1 commit into
Draft
[DO NOT REVIEW] Add MXFP8 recipes for the qwen3 and gpt_oss MoE models#4372anijain2305 wants to merge 1 commit into
anijain2305 wants to merge 1 commit into
Conversation
anijain2305
requested review from
fegin,
tianyu-l,
wconstab and
wwwjn
as code owners
August 28, 2026 22:44
anijain2305
marked this pull request as draft
August 28, 2026 22:44
anijain2305
force-pushed
the
mxfp8-moe-model-configs
branch
from
August 31, 2026 06:00
0d7ac75 to
18027d9
Compare
anijain2305
force-pushed
the
mxfp8-moe-model-configs
branch
from
August 31, 2026 06:08
18027d9 to
89ff3b3
Compare
anijain2305
force-pushed
the
mxfp8-moe-model-configs
branch
from
August 31, 2026 18:24
89ff3b3 to
fc9b778
Compare
anijain2305
force-pushed
the
mxfp8-moe-model-configs
branch
from
August 31, 2026 18:38
fc9b778 to
f61ae82
Compare
anijain2305
force-pushed
the
mxfp8-moe-model-configs
branch
from
August 31, 2026 21:17
f61ae82 to
e9b4aec
Compare
anijain2305
force-pushed
the
mxfp8-moe-model-configs
branch
from
August 31, 2026 21:44
e9b4aec to
d471895
Compare
anijain2305
force-pushed
the
mxfp8-moe-model-configs
branch
from
August 31, 2026 22:15
d471895 to
0f31ed6
Compare
anijain2305
force-pushed
the
mxfp8-moe-model-configs
branch
from
September 1, 2026 05:43
0f31ed6 to
28da47c
Compare
anijain2305
force-pushed
the
mxfp8-moe-model-configs
branch
from
September 1, 2026 05:53
28da47c to
68ea709
Compare
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
force-pushed
the
mxfp8-moe-model-configs
branch
from
September 1, 2026 06:14
68ea709 to
7ec2141
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 CLI • Give Feedback 💬