feat: track MoE models with fused-parameter experts and routers #318 - #360
Open
kohsheen1234 wants to merge 1 commit into
Open
feat: track MoE models with fused-parameter experts and routers #318#360kohsheen1234 wants to merge 1 commit into
kohsheen1234 wants to merge 1 commit into
Conversation
|
|
Author
|
@luciaquirke Can you please look into this :) Would love to hear feedback! |
Collaborator
|
Hello, if your code and PR description are unedited by human hand such that they read as clearly agent-generated then I'm not going to read them. Excepting tests where anything goes. Feedback: code looks bloated |
Closes EleutherAI#318 In transformers 5.x an MoE layer stores its experts as one fused 3D nn.Parameter and its router as a 2D one, so neither is an nn.Linear and discover_targets skips both without saying so. On gpt-oss-20b that leaves 5.8% of the parameters attributed: attention and lm_head only. --track_moe_experts attaches an ExpertLinear per expert projection (model.layers.0.mlp.experts.expert_3.gate_up_proj) and tracks the router in place. gpt-oss-20b reaches 97.2%, Mixtral-8x7B 99.7%. An expert only receives the tokens routed to it, which is ragged across examples, while per-example gradients need [N, S, ...]. Each expert's rows are padded into a rectangular grid and the expert GEMM runs there, so _compute_gradient is untouched: padding rows carry zero activations and receive zero output gradient, contributing nothing to the weight gradient, the bias gradient or any Hessian factor. Routers need a different hook. Every family scores its logits inside its own forward, so register_full_backward_hook reports grad_output[0] is None and the gradient never crosses the module boundary; the collector taps the logits tensor instead. Off by default. Tracking replaces the fused experts' forward with a per-expert loop, which measured 3-11x on collection in CPU eager (GPU unmeasured, and is where losing the fused grouped-matmul kernel will cost most). When tracking is off and a model has fused experts, the collector warns rather than skipping silently, since silence was the original problem. Expansion is reversible and leaves state_dict() untouched, so a model is unaffected once collection ends. Index size grows with layers * experts * 2, ~50 to ~1,600 modules on gpt-oss-20b. --projection_target global absorbs that entirely; --filter_modules narrows it. Tested against per-sample autograd backward passes on gpt-oss, Mixtral, Qwen3-MoE, OLMoE and DeepSeek-V3, spanning both weight orientations, gated and non-gated experts and both router styles, plus a locally declared has_gate=False module for the up_proj path. Also covered: expansion does not change the model's outputs and reverses exactly, works from a model or its base_model, EK-FAC factors use each expert's routed-row mask, and the opt-out wins over an already-expanded model. The CUDA bf16/fp16 tests are written but have not been run. Not covered: llama4 and longcat_flash fuse experts without the use_experts_implementation contract detection relies on; --attribute_tokens is rejected alongside fused experts, since under top-k routing one token feeds several experts and its gradient rows cannot line up with token positions; and --optimizer_state still derives no normalizers for fused expert parameters.
kohsheen1234
force-pushed
the
feat/moe-fused-experts
branch
from
July 29, 2026 12:31
0233235 to
1aa250e
Compare
Collaborator
|
We're moving documentation from the README into the docs so the new MOE info would go in there too, in the .rst format |
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
3 participants
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.
Closes #318
Problem.
transformers5.x stores MoE experts as a fused 3Dnn.Parameterand the router as a bare 2Dnn.Parameter. Neither is annn.Linear, sodiscover_targetsskipped both silently — gpt-oss-20b had 5.8% of its parameters attributed, with no warning.Solution. Behind
--track_moe_experts, each expert projection becomes its own module (...mlp.experts.expert_3.gate_up_proj) and the router is tracked in place: 5.8% → 97.2% on gpt-oss-20b, 3.2% → 99.7% on Mixtral-8x7B. Each expert's ragged routed tokens are padded into an[N, L, ·]grid, so_compute_gradientis unchanged. Off by default, because enabling tracking replaces the fused experts' forward and costs 3–11x on collection (benchmarks/moe_padding_overhead.py); when off, bergson now warns that experts are being skipped instead of staying silent.Done. 32 tests passing — per-example expert and router gradients match per-sample autograd across gpt-oss, Mixtral, Qwen3-MoE, OLMoE and DeepSeek-V3. 6 CUDA bf16/fp16 tests are written but skipped and unverified (no GPU available to me). Not covered:
llama4/longcat_flash(fused params, no decorator contract) and--attribute_tokens.