#10189·ms-swift

[Bug] Qwen3.5 MoE multimodal DPO drops output_router_logits and crashes with aux_loss=None

Author: hkxxxxxCreated Sep 17, 2026Updated Sep 17, 2026

I have searched the existing issues and did not find a report for this multimodal native DPO parameter-passthrough failure.

Environment

  • ms-swift: 4.5.2 (affected local checkout)
  • transformers / trl / torch: exact versions are available from the original environment but are omitted from this sanitized report
  • Python: 3.10.21
  • CUDA: 12.9
  • GPU: NVIDIA H100 80GB
  • Model: Qwen3.6-35B-A3B (Qwen3.5 MoE architecture)
  • Training: native Swift DPO/RLHF + LoRA + DeepSpeed
  • Template: Qwen3.5 multimodal template

Configuration

The relevant option is:

--router_aux_loss_coef 1e-3

When this coefficient is positive, the DPO trainer sets:

python
batch['output_router_logits'] = True

Symptom

The model initializes successfully and NCCL setup completes, but the first training step fails with:

TypeError: unsupported operand type(s) for *: 'float' and 'NoneType'

The failure occurs in the DPO trainer at:

python
losses = losses + self.aux_loss_coef * model_output['aux_loss']

model_output['aux_loss'] is None.

Root cause

The Qwen3.5 multimodal template reuses the Qwen2VL _post_encode path. That path returns a new kwargs dictionary containing inputs_embeds. Template.pre_forward_hook then restores selected values from the original kwargs using an explicit whitelist.

The whitelist does not contain output_router_logits, so the flag set by the DPO trainer is dropped before the Qwen3.5 MoE model forward. The model consequently uses its default output_router_logits=False, leaves aux_loss=None, and the DPO trainer multiplies a float by None.

The effective call path is:

DPO sets output_router_logits=True
  -> multimodal _post_encode rebuilds kwargs
  -> pre_forward_hook drops output_router_logits
  -> Qwen3.5 MoE computes no aux_loss
  -> DPO crashes on float * None

This is distinct from NCCL initialization and from independent CUDA/ABI loading issues that may appear later during distributed shutdown.

Minimal fix

Add output_router_logits to the kwargs restoration whitelist in Template.pre_forward_hook:

python
'output_router_logits',

This preserves the caller-provided flag and allows the model to compute the router auxiliary loss. It does not change the DPO loss formula or convert aux_loss=None to zero.

Validation

With the one-line change, the same DPO run using --router_aux_loss_coef 1e-3 completes the first training step and reports a numeric auxiliary loss, for example:

'nll_loss': '0.2953', 'aux_loss': '8.07'

I am preparing a small regression test that simulates a post-encode path rebuilding kwargs and verifies that output_router_logits=True reaches the model forward.

Questions

  1. Should output_router_logits be added to the generic Template.pre_forward_hook restoration whitelist?
  2. Should a regression test cover a multimodal Qwen3.5 MoE DPO/SFT batch with router_aux_loss_coef > 0?
  3. Should the hook preserve model-forward kwargs more generically instead of maintaining a manually curated whitelist?