modelscope / modelscope/ms-swift

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

Open Beginner friendly
#10,189 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
15.7k
Forks
1.7k
Avg merge
1d 16h
Merged PRs (30d)
136

Description

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:

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:

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:

'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?

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start at Template.pre_forward_hook and trace the Qwen3.5 multimodal _post_encode path, focusing on how rebuilt kwargs are restored. Reproduce the DPO configuration with a positive router_aux_loss_coef, then add the regression coverage described in the issue and verify that output_router_logits reaches model forward and aux_loss is numeric during the first training step.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
2/5
Estimated time
1-3 hours
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
75/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.