Repository navigation
Honor config.output_router_logits in the MoE VLM wrappers - #48885
Conversation
Four `ForConditionalGeneration` wrappers read the flag straight off `kwargs` instead of resolving it against the config, so setting `config.output_router_logits = True` did nothing and only a forward argument turned the auxiliary loss on: Qwen3-VL-MoE, GLM4V-MoE, Qwen3.5-MoE and Qwen4-Exp. Every text-only MoE honors both, Qwen3-MoE being the reference. Resolve the argument against the text config before the backbone call, as `Qwen3MoeForCausalLM` does, and pass it down explicitly. Two modular files cover all four, since Qwen3.5-MoE and Qwen4-Exp inherit the wrapper from Qwen3-VL-MoE. `Qwen3VLMoeTextConfig` had `output_router_logits` deleted in #41250 because nothing read it, and `Glm4vMoeTextConfig` never had it, inheriting from the auxiliary-loss-free `Glm4MoeConfig`. Both need it now, so restore one and add the other.
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
Same gap as the other four, one step further along. `MiniMaxM3VLTextForCausalLM` implements the auxiliary loss in full, while `MiniMaxM3SparseForConditionalGeneration` has no `output_router_logits` argument and no auxiliary-loss block, and neither of its output types carries `aux_loss`. Add the argument and resolve it against the text config like the others, carry `router_logits` through the backbone output instead of dropping it when the output is rebuilt, and add `aux_loss` and `router_logits` to the output types.
| def test_output_router_logits_from_config(self): | ||
| """`config.output_router_logits` turns the auxiliary loss on, and an explicit argument wins over it.""" | ||
| config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() | ||
| model = Glm4vMoeForConditionalGeneration(config).to(torch_device).eval() | ||
| input_ids = inputs_dict["input_ids"] | ||
|
|
||
| assert model(input_ids).aux_loss is None | ||
| assert model(input_ids, output_router_logits=True).aux_loss is not None | ||
|
|
||
| model.config.get_text_config().output_router_logits = True | ||
| assert model(input_ids).aux_loss is not None | ||
| assert model(input_ids, output_router_logits=False).aux_loss is None | ||
|
|
There was a problem hiding this comment.
i thought we had a test in common mixin, but seeing that the PR never got merged due to merge-queue checkd (#46642)
Might be worth to move into mixin which enforces for all models
There was a problem hiding this comment.
ah also the linked PR added aux-loss in gemma4 family, i suppose it is still missing so let's add (if it is still needed)
The modular graph carried the recorder and the config field into models this branch never named: DeepSeek-OCR2 and GLM5-Next through DeepseekV3/Glm4Moe, GLM4V-MoE through Glm4MoeConfig, Youtu through DeepseekV3Config. CI's modular check regenerates everything, so those files have to match. Youtu is dense and DeepSeek-OCR2's wrapper does not carry router logits, so both delete the inherited `output_router_logits` alongside the other MoE fields they already drop. GLM4V-MoE reads the flag as a forward kwarg and #48885 wires it explicitly, so it is allow-listed here until then. Cohere2-MoE's head now takes and resolves the flag like the others, which is also what makes the checker see it.
Cyrilvallez
left a comment
Cyrilvallez
left a comment
There was a problem hiding this comment.
LGTM, let's just make the test general instead of model specific.
We can either merge both PR (the one linked by @zucchini-nlp) or merge separately, no strong opinion.
One important thing to consider is that the general test from the linked PR relies on the config attr to be present to trigger the test, which is exactly what both PR fix -> so would not detect faulty configs. Maybe we should find a better way to detect MoEs to have a more robust test, maybe if "router_logits"is in _can_record_outputs is a more robust way
| def test_output_router_logits_from_config(self): | ||
| """`config.output_router_logits` turns the auxiliary loss on, and an explicit argument wins over it.""" | ||
| config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() | ||
| model = Glm4vMoeForConditionalGeneration(config).to(torch_device).eval() | ||
| input_ids = inputs_dict["input_ids"] | ||
|
|
||
| assert model(input_ids).aux_loss is None | ||
| assert model(input_ids, output_router_logits=True).aux_loss is not None |
There was a problem hiding this comment.
As @zucchini-nlp said, would be better as a general test instead!
…fig' into moe-vlm-output-router-logits-config
|
Moved to It runs on every model whose config has Gemma4: it has no |
zucchini-nlp
left a comment
zucchini-nlp
left a comment
There was a problem hiding this comment.
A few tests failing, can merge after fixing them
| *(getattr(t, "__dataclass_fields__", {}).keys() for t in (get_args(return_type) or (return_type,))) | ||
| ) | ||
| if "router_logits" not in output_fields: | ||
| self.skipTest(f"{model_class.__name__} does not declare router_logits in its output type.") |
There was a problem hiding this comment.
unrelated to this pr - if config has output_router_logits, it should prob have the field in outputs of BaseModels
Or if this is to guard a base model vs generative model, then we can filter all_model_names by get_values(MODEL_MAPPING_NAMES)
…coder text config in the common test
|
[For maintainers] Suggested jobs to run (before merge) run-slow: glm4v_moe, jetmoe, minimax_m3_vl, qwen3_5_moe, qwen3_vl_moe, qwen4_exp |
|
Fixed the 4 failures in 5afac8b: JetMoe ignored run-slow: glm4v_moe, jetmoe, minimax_m3_vl, qwen3_5_moe, qwen3_vl_moe, qwen4_exp |
CI recapDashboard: View test results in Grafana |
…ngface#48920) * Return router logits from the MoE models that balance with a router bias 17 MoE models compute router logits in their gate and drop them: the MoE block unpacks the gate's tuple past them, nothing records them, and the LM head returns a plain `CausalLMOutputWithPast`. So `output_router_logits=True` is silently ignored on DeepSeek-V3/V3.2, GLM4-MoE and Lite, GLM-MoE-DSA, HY-V4, Solar-Open, MiMo-V2-Flash, AXK1/2, LongCat-Flash, Kimi-Linear, dots1, EXAONE-MoE, LFM2-MoE and Nemotron-H. Two of them, EXAONE-MoE and Kimi-Linear, already declared the recorder; the head dropped the value when rebuilding its output field by field. These models balance their experts with a per-expert bias rather than an auxiliary loss, so not computing a loss is right. The logits are still what you need to inspect expert load, and what a trainer needs to compute its own balancing term. Same shape as Qwen3-MoE and AFMoE: a `router_logits` recorder on the pretrained base, the backbone returns `MoeModelOutputWithPast`, the LM head takes `output_router_logits`, resolves it against the config, passes it down and returns `MoeCausalLMOutputWithPast`, and the config gains the field. One change in `deepseek_v3` propagates through the modular graph to eleven of them; `dots1`, `exaone_moe`, `lfm2_moe` and `nemotron_h` get it by hand. `LongcatFlashTopkRouter` now returns its logits first, like the DeepSeek-V3 router it derives from. `aux_loss` stays `None` on all of them, and the docstring on `MoeCausalLMOutputWithPast` now says so instead of tying it to `labels`. A new common test, `test_router_logits_are_returned`, holds every model that records router logits to actually returning them from heads whose output type carries the field. It caught Doge resolving `output_router_logits` and never passing it to its backbone, fixed here too. * Also return router logits from DeepSeek-V2, HunYuan-V1-MoE and Mistral4 huggingface#48911 listed the bias-balanced models. These three have neither a bias nor a loss in HF and were hiding their logits the same way: thin modulars over Llama with a router that already returns `router_logits` first, and heads that return `CausalLMOutputWithPast`. Same recipe as the others. * Return Llama 4's router logits and drop the coefficient nobody chose `Llama4TextConfig` declared `output_router_logits` and `router_aux_loss_coef` and the modeling read neither. The recorder existed but pointed at `Llama4TextMoe`, which returns hidden states; the router is what returns `(router_scores, router_logits)`. Both heads then rebuilt their outputs without the field. Point the recorder at `Llama4Router`, index 1, return `MoeModelOutputWithPast` from the text model and `MoeCausalLMOutputWithPast` from the causal head, and carry `router_logits` through `Llama4CausalLMOutputWithPast`. `output_router_logits=True` now returns one tensor per MoE layer of shape `(batch * seq, num_local_experts)`. `router_aux_loss_coef` goes. Meta published no training recipe and their reference implementation has no balancing loss, so the 0.001 was the porting template's default, not a choice. Downstream code that reads the coefficient as "this model trains with the loss" was crashing on Llama 4 because of it. * Give Cohere2-MoE the output_router_logits config key it was already honouring Its backbone is decorated with merge_with_config_defaults, so the kwarg already worked; only the config default was missing, which made the key an unreliable marker for the one model that returned router logits without declaring it. * Regenerate the downstream models and keep check_config_attributes honest The modular graph carried the recorder and the config field into models this branch never named: DeepSeek-OCR2 and GLM5-Next through DeepseekV3/Glm4Moe, GLM4V-MoE through Glm4MoeConfig, Youtu through DeepseekV3Config. CI's modular check regenerates everything, so those files have to match. Youtu is dense and DeepSeek-OCR2's wrapper does not carry router logits, so both delete the inherited `output_router_logits` alongside the other MoE fields they already drop. GLM4V-MoE reads the flag as a forward kwarg and huggingface#48885 wires it explicitly, so it is allow-listed here until then. Cohere2-MoE's head now takes and resolves the flag like the others, which is also what makes the checker see it. * Return router logits from the Qwen3-Omni-MoE thinker * Inherit the MoE backbones and heads through modular instead of redefining them DeepseekV2Model inherits OlmoeModel, DeepseekV2ForCausalLM is the one router-logits head without a balancing loss, and DeepSeek-V3, HunYuan-V1-MoE, Mistral4, Dots1, Exaone-MoE and LFM2-MoE inherit from those. * Drop duplicate output_router_logits field in Glm4vMoeTextConfig after merge * Address review: rely on the output-capturing decorators, inherit where possible, fix Step3p7 and DeepSeek-OCR2, rework the test * Return router logits from Aria, Inkling and Gemma4 * Align the aux_loss docstring with router_logits * Drop the subTest and the per-model output_router_logits allowances * Address review: keep step3p7 on Gemma3TextModel, no router logits for gemma4_unified, Aria inherits DeepseekV2ForCausalLM
Four
ForConditionalGenerationwrappers read the flag straight offkwargsinstead of resolving it against the config, soconfig.output_router_logits = Truedoes nothing and only a forward argument turns the auxiliary loss on. Every text-only MoE honors both.Qwen3MoeForCausalLMis the reference:transformers/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py
Lines 640 to 642 in 0a959de
transformers/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py
Line 422 in 0a959de
transformers/src/transformers/models/glm4v_moe/modular_glm4v_moe.py
Line 384 in 0a959de
transformers/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
Line 2068 in 0a959de
transformers/src/transformers/models/qwen4_exp/modeling_qwen4_exp.py
Line 2525 in 0a959de
Qwen3.5-MoE and Qwen4-Exp carry both patterns in the same file, the text model resolving against the config and the wrapper ignoring it:
transformers/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
Lines 1897 to 1899 in 0a959de
vs
transformers/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
Line 2068 in 0a959de
After, all three agree with the reference, and an explicit
output_router_logits=Falsestill wins over a configTrue.Changes
Resolve the argument against the text config before the backbone call, the way
Qwen3MoeForCausalLM.forwarddoes, and pass it down explicitly. Two modular files cover all four, sinceQwen3_5MoeForConditionalGeneration(Qwen3VLMoeForConditionalGeneration)andQwen4ExpForConditionalGeneration(Qwen3_5MoeForConditionalGeneration).The two configs also need the field:
Qwen3VLMoeTextConfighadoutput_router_logitsdeleted in 🚨 Validate config attributes #41250, correctly at the time, since nothing read it:transformers/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py
Line 119 in 0a959de
Glm4vMoeTextConfignever had it. It inherits fromGlm4MoeConfig, and GLM4-MoE balances its experts without an auxiliary loss. GLM4V-MoE does compute one, so the field is added next to therouter_aux_loss_coefthe config already carries:transformers/src/transformers/models/glm4v_moe/modular_glm4v_moe.py
Line 109 in 0a959de
One test per modular source, asserting the config default is honored and that the explicit argument wins over it.
Why it matters
Downstream code that decides "can this model produce the auxiliary loss?" from the config gets the wrong answer in both directions. This is what TRL's
router_aux_loss_coefwas keyed on (huggingface/trl#7222): the coefficient was silently ignored on Qwen3-VL-MoE and GLM4V-MoE.Not covered here
This aligns how the flag is read, not what every MoE does with it. Four models still expose
aux_lossin their output type and never populate it, sooutputs.aux_lossisNoneeven with the flag on:output_router_logitsargumentrouter_logitsaux_lossHYV3ForCausalLMis explicit about it:aux_loss=None, # Not used in this model. Either the loss should be computed from the router logits these models already return, or the field should be dropped from their output types. That is a per-architecture question, so it belongs in its own PR.