Llama 4 declares `output_router_logits` and `router_aux_loss_coef` but never reads them

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

Llama4TextConfig declares output_router_logits and router_aux_loss_coef, and modeling_llama4.py never reads either. Zero matches for both in the file. The model returns a plain CausalLMOutputWithPast, so there is no aux_loss and no router_logits to be had:

import torch
from transformers import Llama4ForCausalLM
from transformers.models.llama4.configuration_llama4 import Llama4TextConfig

config = Llama4TextConfig(vocab_size=99, hidden_size=32, intermediate_size=64, intermediate_size_mlp=64,
                          num_hidden_layers=2, num_attention_heads=4, num_key_value_heads=2,
                          num_local_experts=4, num_experts_per_tok=2, head_dim=8)
print("config.output_router_logits :", config.output_router_logits)
print("config.router_aux_loss_coef :", config.router_aux_loss_coef)

model = Llama4ForCausalLM(config).eval()
out = model(input_ids=torch.randint(0, 99, (1, 8)), output_router_logits=True)
print("output type                 :", type(out).__name__)
print("has aux_loss                :", "aux_loss" in type(out).__dataclass_fields__)
print("has router_logits           :", "router_logits" in type(out).__dataclass_fields__)

config.output_router_logits : False
config.router_aux_loss_coef : 0.001
output type                 : CausalLMOutputWithPast
has aux_loss                : False
has router_logits           : False

Same for Llama4ForConditionalGeneration. Both fields have been there since Llama 4 landed in #37307 on 2025-04-05.

The cost is downstream. Code that asks "can this model produce the load-balancing auxiliary loss?" reads the config, gets True for "the field exists", requests it and then crashes. TRL does exactly this, and SFTTrainer(model="trl-internal-testing/tiny-Llama4ForCausalLM") with stock defaults dies:

File "trl/trainer/sft_trainer.py", line 1913, in compute_loss
    aux_loss = outputs.aux_loss
AttributeError: 'CausalLMOutputWithPast' object has no attribute 'aux_loss'

Expected behavior

Either implement the auxiliary loss, or drop the two fields, in the spirit of #41250. I'd guess drop, since nothing reads them, but Llama 4's own training recipe may argue for implementing. Happy to send whichever PR you prefer.

Related: #48885 wires output_router_logits through the MoE VLM wrappers, #48886 covers four models that expose aux_loss and never set it. Same theme, different manifestations.