router_aux_loss_coef is silently ignored on MoE configs without output_router_logits
Reproduction
router_aux_loss_coef silently does nothing on half of the MoE architectures, because whether a model counts as a MoE is inferred from a config field many MoE configs do not declare:
# trl/trainer/sft_trainer.py
is_moe = getattr(text_config, "output_router_logits", None) is not None
self.aux_loss_enabled = is_moe and self.args.router_aux_loss_coef != 0.0from datasets import Dataset
from trl import SFTConfig, SFTTrainer
dataset = Dataset.from_dict({"prompt": ["What is 2+2?"] * 8, "completion": [" Four."] * 8})
for model_id in ("hf-internal-testing/tiny-random-Qwen3MoeForCausalLM",
"hf-internal-testing/tiny-random-DeepseekV3ForCausalLM"):
trainer = SFTTrainer(
model=model_id,
args=SFTConfig(output_dir="out", max_steps=2, logging_steps=1, report_to=[],
router_aux_loss_coef=0.001),
train_dataset=dataset,
)
trainer.train()
keys = {k for h in trainer.state.log_history for k in h}
print(f"{model_id}: aux_loss logged = {'aux_loss' in keys}")Outcome
tiny-random-Qwen3MoeForCausalLM: aux_loss logged = True
tiny-random-DeepseekV3ForCausalLM: aux_loss logged = FalseBoth are MoE, both were asked for the same coefficient, one gets the term and the other does not. No warning either way.
On transformers main, 34 MoE configs declare output_router_logits and 30 do not. The second group includes DeepSeek-V2/V3/V3.2, GLM4-MoE, GLM4V-MoE, Qwen3-VL-MoE, Mistral4, dots1, Kimi-Linear and Step3p7. For all of those, setting router_aux_loss_coef is a no-op.
Expected behavior
Either detect MoE from something every MoE has, e.g. num_experts / num_local_experts on the text config, or keep the current check and warn when router_aux_loss_coef is non-zero but the term cannot be applied. Silently ignoring the argument is the part to avoid: the run looks configured and is not.
System info
trl 1.13.0.dev0, transformers 5.18.0.dev0, peft 0.21.0, accelerate 1.15.0, torch 2.13.0, one H100.
This issue is not open to external contribs
Source: huggingface/trl