四个 MoE 模型在其输出类型中暴露了 `aux_loss`，并且从未设置它

Who can help? @ArthurZucker @Cyrilvallez @vasqu ### Information - [x] The official example scripts - [ ] My own modified scripts ### Tasks - [x] 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