Qwen3 MOE models w/non-empty `mlp_only_layers` fail when `output_router_logits=True`

[ntenenz]

System Info

MLPs insert a None into the router logit list. Therefore, when enabled in the model, they need to be filtered out either before or within the aux_loss function.

Who can help?

@ArthurZucker

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

from transformers import AutoConfig, AutoModelForCausalLM

config = AutoConfig.from_pretrained("Qwen/Qwen3-30B-A3B")  # or any qwen3_moe model
config.update({"mlp_only_layers": [0]})  # or any non-empty list
model = AutoModelForCausalLM.from_config(config)

_ = model(INPUT_TOKENS, output_router_logits=True)  # raises an error in the aux_loss function

Expected behavior

Model should output router logits and not raise an exception.

[SwiftAkira]

I’d like to work on this, can you confirm filtering out None router logits is the right approach?

[SwiftAkira]

Hi all,
This issue is now resolved by [PR #39206](#39206).

• The fix ensures that Qwen3 MoE models with non-empty mlp_only_layers no longer include None values in the router logits when output_router_logits=True.
• The router logits are now filtered so only valid tensors from MoE layers are included, matching the Qwen2 MoE implementation pattern.
• Auxiliary loss computation now works correctly in all configurations, including edge cases where all layers are MLP-only.

I’ve run comprehensive tests on all combinations of mlp_only_layers and confirmed the error no longer occurs.

Reference:
PR #39206: [fix: filter None router logits in Qwen3 MoE and handle empty router logits (#39203)](#39206)

cc @ArthurZucker @ntenenz

[ArthurZucker]

Having a look! 🤗

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.