modelscope / modelscope/ms-swift
[Bug] Qwen3.5 MoE multimodal DPO drops output_router_logits and crashes with aux_loss=None
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
- Should
output_router_logitsbe added to the genericTemplate.pre_forward_hookrestoration whitelist? - Should a regression test cover a multimodal Qwen3.5 MoE DPO/SFT batch with
router_aux_loss_coef > 0? - Should the hook preserve model-forward kwargs more generically instead of maintaining a manually curated whitelist?
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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