Add Qwen3 MoE support (M0-M7): ops, autograd, config, weight loading, model forward, real-checkpoint gate - #1
Open
Billy1900 wants to merge 9 commits into
Conversation
Design-only milestone: define moe_router_top_k, moe_expert_ffn, and moe_combine shape/dtype/oracle contracts in docs/OPERATOR_CONTRACTS.zh-CN.md before any ops.h, coverage_manifest, or kernel code lands, per the repo's correctness-first operator workflow. First version deliberately computes all experts and masks non-selected ones (no gather/dispatch yet). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Implement the naive CPU float32 reference for moe_router_top_k,
moe_expert_ffn, and moe_combine per the M0 contracts: softmax+top-k
routing, compute-all-experts-then-mask SwiGLU FFN (O(num_experts),
not O(k)), and weighted combine. Each explicitly rejects a HIP Tensor
("no HIP kernel yet") rather than silently falling back.
Adds hand-computed CPU unit tests (tests/ops/ops_test.cpp), wires the
three-way PyTorch oracle (tests/torch/operator_oracle.cpp +
python/tests/test_operator_parity.py), and registers all three ops in
tests/coverage_manifest.json so scripts/audit_test_coverage.py passes.
Verified: 288/288 microllm_tests pass on a CPU-only
(MICROLLM_ENABLE_HIP=OFF) build, including the 3 new MoE cases with
hand-computable expected values, and the coverage audit passes. The
PyTorch-side oracle (TorchOps.OperatorParity) could not be executed in
this environment (no PyTorch/libtorch install) — both sides were
written to compute the identical formula and should be run once on a
machine with PyTorch before this milestone is considered closed; see
docs/development/2026-09-04-m1-qwen3-moe-cpu-reference.md.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Point CMAKE_PREFIX_PATH at a local conda env's CUDA-built PyTorch (torch==2.11.0+cu130) to build microllm_operator_oracle and run TorchOps.OperatorParity for real, instead of only reasoning that both oracle sides compute the same formula. All 4 parity tests pass, including the forward-value comparison for moe_router_top_k, moe_expert_ffn, and moe_combine against independently written PyTorch references. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Implement moe_router_top_k/moe_expert_ffn/moe_combine HIP kernels in src/ops/hip/basic_kernels.hip, mirroring the M1 CPU reference exactly: router selects top-k via an O(1)-state repeated-scan (no per-token array sized by num_experts, so no artificial expert-count cap), the expert FFN is split into a hidden-projection kernel and a down- projection-plus-mask kernel through a caller-allocated workspace, and combine reads only the k selected experts. Wire real dispatch into the ops.cpp device().is_hip() branches, replacing the M1 placeholder throws. Important: this machine has no ROCm/HIP toolchain or AMD GPU, so the .hip kernel code has never been compiled by hipcc or run on hardware -- only hand-reviewed against this file's existing kernel idioms (the argmax tie-break comparator, the embedding NaN-on-out-of-range pattern). The CPU-facing dispatch code in ops.cpp is compiler-verified (MICROLLM_ENABLE_HIP=OFF build, 55/55 CpuOpsTest pass, unchanged from M1). Per agreement with the user, the three-way CPU/HIP/PyTorch oracle in tests/ops/hip_ops_test.cpp is deferred until AMD hardware is available; see docs/development/2026-09-04-m2-qwen3-moe-hip-kernels.md for the full accounting of what is and isn't verified. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Add CPU backward primitives (moe_router_top_k_backward,
moe_expert_ffn_backward returning a new TensorQuad, moe_combine_backward)
and autograd::Value-level wiring (moe_router_top_k returning a new
MoeRouterResult{Tensor indices; Value weights}, moe_expert_ffn,
moe_combine), each gated by PyTorch autograd over the forward rather
than a hand-derived formula.
Key correctness points: expert_indices is always a plain Tensor, never
a Value, matching embedding()'s convention -- top-k selection is not
differentiable. The router's logit gradient is dense (softmax couples
every expert through its shared denominator), not sparse; only the
selection itself is non-differentiable. The expert FFN's per-expert
weight gradients are exactly zero for any expert a token didn't
select -- a natural consequence of the forward's mask-multiply, not a
hand-written scatter -- while combine's expert_output gradient is a
genuine scatter-add.
Verified three ways: CPU finite-difference/hand-value unit tests
(tests/ops/ops_test.cpp), a full-graph CPU test chaining all three ops
through one backward() call (tests/autograd/autograd_test.cpp), and
the actual PyTorch parity suite (TorchOps.OperatorParity, all 4 tests
pass) built against a CUDA PyTorch from a local conda env -- the same
technique used to close the M1 gap. 292/292 microllm_tests pass on a
CPU-only build; coverage audit passes (tensor_ops 202->205, graph_api
45->48).
CPU only: HIP backward throws explicitly, deferred until M2's forward
HIP kernels get their first real hardware verification (see
docs/development/2026-09-04-m3-qwen3-moe-autograd-backward.md).
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Accept model_type=qwen3_moe in load_huggingface_config, parsing num_experts, num_experts_per_tok, norm_topk_prob, and moe_intermediate_size into four new ModelConfig fields (moe_num_experts/moe_num_experts_per_tok/moe_intermediate_size/ moe_norm_topk_prob). validate() enforces an all-zero-or-consistent invariant across them. Explicitly rejects decoder_sparse_step != 1, a non-empty mlp_only_layers (both mean unsupported per-layer dense/MoE mixing), and any router_aux_loss_coef (an unimplemented training-time loss) rather than silently ignoring them. parameter_count()/weight_bytes() explicitly throw for MoE configs: the exact per-expert tensor layout is a weight-loading decision (M5), and a wrong dense-only count would be worse than refusing to answer. The assumed MoE field names follow Hugging Face's Qwen2MoeConfig convention; this machine has no network access to verify them against a real Qwen3-MoE config.json, so this is parsing-only against synthetic fixtures for now -- flagged explicitly in docs/development/2026-09-04-m4-qwen3-moe-config-parsing.md for re-checking once a real checkpoint is available (M7). 295/295 microllm_tests pass on a CPU-only build (+3 new); coverage audit unaffected (config.h/huggingface.h aren't part of its audited surface). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Give TransformerModel real per-expert MoE parameters so the existing
strict/non-strict load_state_dict machinery works end to end, per an
explicit scope decision made with the user (full integration into the
4167-line model.cpp, not a decoupled module deferred to M6).
Add MoeFeedForward: one router Linear plus per-expert gate/up/down
Linear vectors, registered as real named parameters
(blocks.N.moe.router.weight, blocks.N.moe.experts.E.{gate,up,down}_proj
.weight). Block now holds std::optional<FeedForward>/
std::optional<MoeFeedForward> selected by an is_moe_ flag; all 18
former feed_forward_.xxx() call sites became is_moe_ branches.
qwen_style_weight_mapping gained a matching MoE branch mapping to an
assumed Hugging Face external naming (unverified against a real
Qwen3-MoE checkpoint -- flagged in the dev record, same caveat as M4).
Forward and every one-way inference preparation (BF16/FP8/INT8, the
BF16 FFN Arena) stay explicitly unimplemented -- each throws rather
than silently running a wrong or partial computation; that's M6. Fixed
two related gaps: the model constructor's parameter-count cross-check
now skips MoE configs (ModelConfig::parameter_count() already throws
for them since M4), and prepare_bf16_ffn_inference gained an explicit
MoE guard instead of relying on its name-filter accidentally matching
zero MoE tensors.
Verified: 298/298 microllm_tests pass on both a CPU-only build and
under ASan/UBSan (was 295 after M4; +3 new tests covering the exact
tensor-count assertion, the strict-mode missing/unexpected/incompatible
triple, and the explicit forward/preparation rejections). No dense-model
regression -- the is_moe_ branch is inert for every pre-existing config.
Coverage audit unaffected. Full accounting in
docs/development/2026-09-04-m5-qwen3-moe-weight-loading.md.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Give TransformerModel a real MoE forward, gated against HF Transformers' actual, unmodified Qwen3MoeSparseMoeBlock.forward() (transformers 5.8.0, via a local conda env) at atol=2e-3/rtol=2e-3 -- not a hand-rolled equivalent, per this milestone's standard. Before writing any M6 code, checked that same transformers install against M5's design and found M5's internal MoE weight layout was wrong: it assumed per-expert separate gate/up/down Linears, but the real Qwen3MoeExperts module stores exactly two packed parameters per layer regardless of expert count -- gate_up_proj (gate and up fused, [num_experts, 2*ffn_dim, dim]) and down_proj. Router math and all M4 config field names were already correct. Surfaced this to the user before proceeding; they chose to fix M5's internal representation to match the real checkpoint (keeping ops::moe_expert_ffn's existing separate-gate/up interface) over reopening M1-M3's ops. Changes: MoeFeedForward's gate_up_proj_/down_proj_ are now plain Values in HF's exact native layout, so qwen_style_weight_mapping stays a plain Identity copy -- no new per-expert-slice loading capability needed. Added ops::moe_split_gate_up/_backward (CPU only) and autograd::moe_split_gate_up as a checkpoint-format adapter -- not a routing primitive -- to convert to moe_expert_ffn's expected layout at forward time; down_proj needs only the existing autograd::transpose. MoeFeedForward::forward()/forward_tensor() now actually compute router -> moe_split_gate_up -> moe_expert_ffn -> moe_combine instead of throwing. BF16/FP8/INT8 MoE support remains explicitly unimplemented. Verified: hand/finite-difference tests for the new adapter op; model tests rewritten for the corrected tensor layout (exact-count assertion, strict-mode triple, and a new end-to-end forward+backward-trains-every- parameter test); and the actual PyTorch parity suite, which now includes a deterministic-fixture comparison against the real Qwen3MoeSparseMoeBlock (output + every gradient). 301/301 microllm_tests pass on both a CPU-only build and under ASan/UBSan; coverage audit passes. Full accounting in docs/development/2026-09-04-m6-qwen3-moe-model-forward.md. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This machine turns out to have network access (never re-checked after M4 assumed otherwise). Used it to download a real Qwen3-MoE checkpoint before writing M7's fixture code, and found M6's core decision was wrong: M6 read only transformers' current in-memory Qwen3MoeExperts source (a fused gate_up_proj nn.Parameter) and assumed that matched what's on disk. It doesn't. Both a downloaded tiny checkpoint and the official Qwen/Qwen3-30B-A3B's safetensors index store per-expert separate gate_proj/up_proj/down_proj tensors -- M5's original design. transformers converts these into its packed in-memory shape via cat([gate,up],dim=0) at load time; this repo never needed that packed representation in the first place. Surfaced to the user before reverting, since it meant undoing M6's commit, not just extending it; they chose to revert to per-expert storage over teaching WeightMapping a new many-tensors-into-one-slice loader capability. Revert: MoeFeedForward is back to per-expert std::vector<Linear>. Deleted ops::moe_split_gate_up/_backward and autograd::moe_split_gate_up (solved a problem no real checkpoint has); added ops::moe_stack_experts/_backward_one and autograd::moe_stack_experts (stacks N per-expert tensors into moe_expert_ffn's packed shape, no transpose needed since this repo's own Linear layout already matches). qwen_style_weight_mapping's MoE branch is back to a per-expert Transpose2D loop. Also found and fixed three more real-checkpoint gaps in apps/hf_infer.cpp and config parsing, none of them MoE-routing bugs -- all four were "this code path was never exercised against an actual downloaded, tied, newer-schema checkpoint before": - router_aux_loss_coef's mere presence was rejected outright (M4); every real Qwen3-MoE config serializes it with its default, making every real checkpoint unloadable. Now accepted without effect -- it configures a training loss this repo does not implement. - Newer configs serialize "dtype" instead of "torch_dtype"; both keys now accepted (purely informational metadata either way). - qwen3_tied_weight_aliases assumed every tied checkpoint redundantly serializes lm_head.weight; hf_infer.cpp now checks the weight file before applying that alias instead of applying it unconditionally. - external.model.weight_bytes() (used only for a benchmark metric) throws for MoE the same way parameter_count() deliberately does; replaced with model.parameter_count() * sizeof(float), the model's own live count, for both dense and MoE. Added a fixture entry to data/model_fixtures.toml (amd-quark/tiny-random-qwen3_moe -- a structurally real but randomly-initialized checkpoint; the smallest genuinely pretrained Qwen3-MoE model is far too large to run here). It has no declared license on Hugging Face, same as every comparable tiny-random test repo checked; flagged to the user, who chose to add it anyway with an explicit "unspecified" license note rather than a real open-source license string, since the weights are random (no pretrained content to license) and are never committed to this repository regardless. The actual gate: microllm_hf_infer's greedy-decoded tokens for a fixed prompt exactly match transformers' golden generation on the same checkpoint. tools/prepare_hf_fixture.py prepare/validate both pass. 302/302 microllm_tests pass on both a CPU-only build and under ASan/UBSan; coverage audit passes; the PyTorch parity suite (rebuilt against transformers 5.8.0) passes all 4 tests including the rewritten per-expert MoE model-level gate against the real Qwen3MoeSparseMoeBlock. Existing hf_infer CLI/CTest contracts re-verified against the rebuilt binary. Full accounting in docs/development/2026-09-05-m7-qwen3-moe-real-checkpoint.md. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
There was a problem hiding this comment.
🟡 Changes recommended
The new MoE ops include missing same-device validation (can cause incorrect behavior/UB on HIP) and there’s at least one now-misleading public header comment that should be corrected before approval.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
This PR adds end-to-end Qwen3 MoE support across the operator layer, autograd graph, Hugging Face config parsing, weight loading, and model forward execution, with extensive parity gates (including a real-checkpoint golden-output gate).
Changes:
- Introduces MoE ops (
moe_router_top_k,moe_expert_ffn,moe_combine,moe_stack_experts) plus CPU backward primitives and autograd wiring. - Integrates MoE into
TransformerModel(config validation, HF mapping, per-expert parameters, forward path, and explicit non-support for advanced inference prep). - Adds/extends parity tests, invalid-shape coverage, and development records + fixture registry entry for a real MoE checkpoint gate.
File summaries
| File | Description |
|---|---|
| tests/torch/operator_oracle.cpp | Adds MoE forward/backward oracle emissions and a model-level MoE gate case. |
| tests/ops/ops_test.cpp | Adds CPU unit tests for MoE ops and their backward behavior. |
| tests/model/weights_test.cpp | Adds MoE weight mapping/state_dict strictness tests and forward/grad coverage. |
| tests/model/config_test.cpp | Adds ModelConfig and HF config parsing tests for qwen3_moe. |
| tests/coverage_manifest.json | Registers MoE ops for test coverage auditing. |
| tests/autograd/autograd_test.cpp | Adds autograd graph tests for MoE routing/FFN/combine and stacking. |
| src/ops/ops.cpp | Implements MoE ops (CPU + HIP forward dispatch) and CPU backward primitives. |
| src/ops/hip/kernels.h | Declares HIP kernel launchers for MoE forward ops. |
| src/ops/hip/basic_kernels.hip | Adds readable HIP kernels for MoE forward ops. |
| src/model/model.cpp | Implements MoE FFN module, integrates MoE into Block/TransformerModel forward paths. |
| src/model/huggingface.cpp | Extends HF config loader to accept/validate model_type=qwen3_moe and related fields. |
| src/model/config.cpp | Adds MoE validation rules and makes parameter_count reject MoE configs. |
| src/autograd/autograd.cpp | Adds autograd operations for MoE ops, including mixed Tensor/Value router result. |
| scripts/audit_test_coverage.py | Updates audited public API symbol lists for new MoE-related types. |
| python/tests/test_operator_parity.py | Extends PyTorch parity references and adds an HF Qwen3 MoE block oracle gate. |
| include/microllm/ops/ops.h | Declares MoE ops + backward APIs and introduces TensorQuad. |
| include/microllm/model/huggingface.h | Updates HF config parsing documentation for qwen3_moe. |
| include/microllm/model/config.h | Adds MoE fields to ModelConfig. |
| include/microllm/autograd/autograd.h | Declares MoE autograd APIs and MoeRouterResult. |
| docs/OPERATOR_CONTRACTS.zh-CN.md | Documents MoE routing contracts, milestones, and constraints. |
| docs/development/README.md | Adds an index of the M0–M7 MoE development records. |
| docs/development/2026-09-05-m7-qwen3-moe-real-checkpoint.md | Records the M7 real-checkpoint findings and verification steps. |
| docs/development/2026-09-04-m6-qwen3-moe-model-forward.md | Records M6 model-forward gate design/verification (and historical context). |
| docs/development/2026-09-04-m5-qwen3-moe-weight-loading.md | Records M5 weight loading integration and strictness checks. |
| docs/development/2026-09-04-m4-qwen3-moe-config-parsing.md | Records M4 HF config parsing decisions and validation behavior. |
| docs/development/2026-09-04-m3-qwen3-moe-autograd-backward.md | Records M3 backward primitive design and PyTorch-autograd gating. |
| docs/development/2026-09-04-m2-qwen3-moe-hip-kernels.md | Records M2 readable HIP kernel design and “unverified on hardware” boundary. |
| docs/development/2026-09-04-m1-qwen3-moe-cpu-reference.md | Records M1 CPU reference implementation and oracle verification approach. |
| docs/development/2026-09-04-m0-qwen3-moe-op-contracts.md | Records the initial operator contract design milestone (M0). |
| data/model_fixtures.toml | Adds a pinned tiny-random Qwen3 MoE fixture entry and expected tokens. |
| apps/hf_infer.cpp | Fixes tied-weight alias application and weight-bytes reporting to work with MoE checkpoints. |
Review details
- Files reviewed: 31/31 changed files
- Comments generated: 3
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+2459
to
+2461
| if (expert_indices.dtype() != DType::Int32) { | ||
| throw std::invalid_argument("moe_expert_ffn expert_indices must be an int32 tensor"); | ||
| } |
Comment on lines
+2655
to
+2659
| const auto& expert = experts[index]; | ||
| if (expert.dtype() != DType::Float32 || expert.shape() != first.shape()) { | ||
| throw std::invalid_argument( | ||
| "moe_stack_experts requires every expert to share float32 dtype and shape"); | ||
| } |
Comment on lines
+19
to
+22
| // Dense Qwen2/Qwen2.5 and Qwen3 decoder config compatibility, plus Qwen3 MoE | ||
| // (model_type=qwen3_moe) config parsing. Only uniform per-layer MoE | ||
| // (decoder_sparse_step=1, empty mlp_only_layers) is supported; MoE weight | ||
| // loading and model-level forward are not implemented by this milestone. |
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.
Summary
Adds Qwen3 MoE support to the ops/autograd/model layers, following a nine-commit milestone plan (M0 through M7):
moe_router_top_k,moe_expert_ffn,moe_combine(design doc only, no code)model_type=qwen3_moeTransformerModel(realstate_dict/load_state_dict, strict-mode missing/unexpected/incompatible triple)TransformerModel::forward()MoE computation, gated against the actual, unmodifiedtransformersQwen3MoeSparseMoeBlock.forward()microllm_hf_infer's greedy-decoded tokens exactly matchtransformers' golden generation on the same weightsNotable course-corrections (both surfaced to the user before proceeding)
gate_up_projparameter) was based on reading onlytransformers' current in-memory module source. Downloading a real checkpoint and inspecting the officialQwen/Qwen3-30B-A3Bsafetensors index in M7 showed real checkpoints actually store per-expert separategate_proj/up_proj/down_projtensors — M5's original design. M6's fused-layout ops (moe_split_gate_up) were deleted and replaced withmoe_stack_experts, which assembles per-expert tensors into the packed shape the routing ops expect.hf_infer(arouter_aux_loss_coefover-rejection, atorch_dtype/dtypeconfig key rename, and a tied-embedding alias that assumed a redundantlm_head.weighttensor always exists on disk).Verification
microllm_testspass on a CPU-only build (MICROLLM_ENABLE_HIP=OFF) and under ASan/UBSanscripts/audit_test_coverage.pypassesTorchOps.OperatorParity(three-way CPU/PyTorch oracle) passes, run against a local PyTorch install including the realtransformersQwen3MoeSparseMoeBlockcomparisonmicrollm_hf_infer's output onamd-quark/tiny-random-qwen3_moeexactly matchestransformers' own greedy generationdocs/development/2026-09-04-m2-qwen3-moe-hip-kernels.mdfor the exact boundary of what is and isn't compiled/testedTest plan
MICROLLM_ENABLE_HIP=ONon a real ROCm machine and runtests/ops/hip_ops_test.cpp(does not yet exist — HIP-side testing deferred until hardware is available, per the M2 record)MICROLLM_ENABLE_HIP=OFFbuild + fullmicrollm_testssuite (CPU + ASan/UBSan)scripts/audit_test_coverage.pyTorchOps.OperatorParitymicrollm_hf_inferagainst a real downloaded Qwen3-MoE checkpointSee
docs/development/2026-09-04-m0-qwen3-moe-op-contracts.mdthroughdocs/development/2026-09-05-m7-qwen3-moe-real-checkpoint.mdfor the full per-milestone record, including every course-correction and its reasoning.🤖 Generated with Claude Code