Four MoE models expose `aux_loss` in their output type and never set it

System Info

• transformers version: 5.17.0.dev0
• Platform: Linux-6.8.0-1057-aws-x86_64-with-glibc2.35
• Python version: 3.13.13
• Huggingface_hub version: 1.29.0
• Accelerate version: 1.14.0
• PyTorch version: 2.13.0+cu130 (GPU: True)

Who can help?

@ArthurZucker @Cyrilvallez @vasqu

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder
• My own task or dataset

Reproduction

Four MoE models return MoeCausalLMOutputWithPast, so aux_loss is part of their public output type, but none of them ever sets it. outputs.aux_loss is None even with output_router_logits=True:

import torch
from transformers import AutoConfig, AutoModelForCausalLM

TINY = dict(hidden_size=32, intermediate_size=64, num_hidden_layers=2, num_attention_heads=4, num_key_value_heads=2, vocab_size=99, num_experts=4, num_experts_per_tok=2)

for model_type in ("qwen3_moe", "afmoe", "cohere2_moe", "hy_v3", "zaya"):
    kwargs = dict(TINY, num_experts_per_tok=1) if model_type == "zaya" else TINY
    config = AutoConfig.for_model(model_type, **kwargs)
    model = AutoModelForCausalLM.from_config(config).eval()
    out = model(input_ids=torch.randint(0, 99, (1, 8)), output_router_logits=True)
    print(f"{model_type:14s} router_logits={out.router_logits is not None}  aux_loss={out.aux_loss}")

qwen3_moe      router_logits=True  aux_loss=2.001878023147583
afmoe          router_logits=True  aux_loss=None
cohere2_moe    router_logits=True  aux_loss=None
hy_v3          router_logits=True  aux_loss=None
zaya           router_logits=True  aux_loss=None

HYV3ForCausalLM says so in the code: aux_loss=None, # Not used in this model.

Three of the four look deliberate. They balance their experts with a per-expert bias added to the routing scores before the top-k, the aux-loss-free scheme, so there is no auxiliary loss to return:

	balancer
AFMoE	expert_bias, torch.topk(scores + expert_bias, ...)
HY-V3	e_score_correction_bias
Zaya	balancing_biases, router_probs.detach() + self.balancing_biases
Cohere2-MoE	none that I can find, plain top-k on the router logits

AFMoE, HY-V3 and Zaya also declare output_router_logits in their configs, which reads as "this model supports the auxiliary loss" when it doesn't.

Cohere2-MoE is the odd one. It has neither an auxiliary loss nor a bias, so nothing balances its experts during training. I don't know whether that matches Cohere's own recipe.

Expected behavior

Two things worth deciding, and I'm happy to send the PR once you pick:

1. A model that never produces an auxiliary loss shouldn't advertise aux_loss in its output type. There's no causal-LM output class today with router_logits but no aux_loss, so this needs a new one, or per-model dataclasses like the VLMs have.
2. Cohere2-MoE either wants the auxiliary loss implemented, or it is bias-free by design and just inherits the class like the other three.

Why it matters beyond tidiness: downstream code can't tell from the outside which MoE can produce the term. The config doesn't say, which is #48885, and the output type doesn't say either, which is this. TRL keys its router_aux_loss_coef on exactly this question (huggingface/trl#7222) and has to guard on outputs.aux_loss is not None at runtime instead.