Qwen3.5: DeepSpeed ZeRO-3 fails to load weights for language_model

[debOliveira]

System Info

• transformers version: 5.4.0
• Platform: Linux (H200 x4)
• Python version: 3.12.0
• DeepSpeed version: 0.18.5
• PyTorch version: 2.8.0+cu128 (CUDA)

Problem

When loading Qwen/Qwen3.5-27B (also tested with 9B) with DeepSpeed ZeRO-3, language_model parameters are reported as MISSING in the load report.

Key                                                                  | Status  | Details
---------------------------------------------------------------------+---------+--------
model.language_model.layers.{0...63}.post_attention_layernorm.weight | MISSING |        
model.language_model.layers.{0...62}.linear_attn.out_proj.weight     | MISSING |        
model.language_model.norm.weight                                     | MISSING |        
model.language_model.layers.{0...62}.linear_attn.in_proj_qkv.weight  | MISSING |        
model.language_model.layers.{0...62}.linear_attn.conv1d.weight       | MISSING |        
model.language_model.layers.{3...63}.self_attn.q_norm.weight         | MISSING |        
model.language_model.layers.{0...62}.linear_attn.dt_bias             | MISSING |        
model.language_model.layers.{0...62}.linear_attn.A_log               | MISSING |        
model.language_model.layers.{0...62}.linear_attn.norm.weight         | MISSING |        
model.language_model.layers.{0...63}.mlp.gate_proj.weight            | MISSING |        
model.language_model.layers.{0...63}.mlp.down_proj.weight            | MISSING |        
model.language_model.layers.{0...62}.linear_attn.in_proj_a.weight    | MISSING |        
model.language_model.layers.{3...63}.self_attn.q_proj.weight         | MISSING |        
model.language_model.layers.{3...63}.self_attn.k_proj.weight         | MISSING |        
model.language_model.layers.{0...63}.input_layernorm.weight          | MISSING |        
model.language_model.layers.{0...63}.mlp.up_proj.weight              | MISSING |        
model.language_model.layers.{0...62}.linear_attn.in_proj_b.weight    | MISSING |        
model.language_model.layers.{0...62}.linear_attn.in_proj_z.weight    | MISSING |        
model.language_model.layers.{3...63}.self_attn.k_norm.weight         | MISSING |        
model.language_model.layers.{3...63}.self_attn.o_proj.weight         | MISSING |        
model.language_model.layers.{3...63}.self_attn.v_proj.weight         | MISSING |        
model.language_model.embed_tokens.weight                             | MISSING | 

Cause hypothesis

In conversion_mapping.py, the language_model weight keys are remapped to model only. This conversion is called in _load_pretrained_model when DeepSpeed ZeRO-3 is turned on.

transformers/src/transformers/conversion_mapping.py

Lines 155 to 157 in d081c71

	"qwen3_5_text": [
	WeightRenaming(source_patterns=r"^model.language_model", target_patterns="model"),
	],

The problem disappears when setting target_patterns="model.language_model" or using ZeRO-2.

Reproduction

from transformers import Qwen3_5ForConditionalGeneration
from transformers.integrations import HfDeepSpeedConfig

ds_cfg = {
    "bf16": {"enabled": True},
    "train_micro_batch_size_per_gpu": 1,
    "train_batch_size": 1,
    "zero_optimization": {
        "stage": 3,
    }
}
dschf = HfDeepSpeedConfig(ds_cfg)
model = Qwen3_5ForConditionalGeneration.from_pretrained("Qwen/Qwen3.5-27B", torch_dtype="auto")

Expected Behavior

Model weights should load correctly with DeepSpeed ZeRO-3.

Related issues

• [BUG] transformers>=5.4.0, Qwen3.5 Moe from_pretrained error #45310
• [Regression] Qwen3.5 saved checkpoint is not correct with save_pretrained API since version 5.4.0 #45216

[zucchini-nlp]

Duplicate of #45216 and will be fixed soon in #45314

[debOliveira]

I tried the recent patch up to commit 9d840ea, but the issue still remains. Note that it works fine on stage 2 of deepspeed, it only raises on stage 3.

Using the same reproduction section above, I get:

Key                                                                  | Status  | 
---------------------------------------------------------------------+---------+-
model.language_model.layers.{0...63}.input_layernorm.weight          | MISSING | 
model.language_model.layers.{0...62}.linear_attn.out_proj.weight     | MISSING | 
model.language_model.layers.{0...62}.linear_attn.dt_bias             | MISSING | 
model.language_model.layers.{0...62}.linear_attn.in_proj_qkv.weight  | MISSING | 
model.language_model.layers.{0...63}.mlp.down_proj.weight            | MISSING | 
model.language_model.layers.{0...62}.linear_attn.in_proj_a.weight    | MISSING | 
model.language_model.layers.{0...63}.post_attention_layernorm.weight | MISSING | 
model.language_model.layers.{0...62}.linear_attn.norm.weight         | MISSING | 
model.language_model.layers.{0...62}.linear_attn.A_log               | MISSING | 
model.language_model.layers.{0...63}.mlp.gate_proj.weight            | MISSING | 
model.language_model.layers.{3...63}.self_attn.v_proj.weight         | MISSING | 
model.language_model.layers.{0...62}.linear_attn.in_proj_z.weight    | MISSING | 
model.language_model.layers.{3...63}.self_attn.o_proj.weight         | MISSING | 
model.language_model.layers.{0...62}.linear_attn.in_proj_b.weight    | MISSING | 
model.language_model.layers.{3...63}.self_attn.q_norm.weight         | MISSING | 
model.language_model.layers.{3...63}.self_attn.q_proj.weight         | MISSING | 
model.language_model.layers.{0...62}.linear_attn.conv1d.weight       | MISSING | 
model.language_model.embed_tokens.weight                             | MISSING | 
model.language_model.layers.{0...63}.mlp.up_proj.weight              | MISSING | 
model.language_model.layers.{3...63}.self_attn.k_proj.weight         | MISSING | 
model.language_model.layers.{3...63}.self_attn.k_norm.weight         | MISSING | 
model.language_model.norm.weight                                     | MISSING | 

[zucchini-nlp]

@Cyrilvallez , then there is definitely smth going on internally in deepspeed

[debOliveira]

I tested the new patch at #45358, but still no success. The weights under model.language_model.layers are being remapped to only model.layers. From what I understand the target being r"^model.(?!(?:language_model.|visual.))" excludes language_model string.

A bad fix that I tested (and works) is

WeightRenaming(
    source_patterns=r"^model.language_model.",
    target_patterns=r"^model.language_model."
)

[BartekKruczek]

Hello 👋

The same problem (PEFT + LoRA + Accelerate + Deepspeed3) on 4 x H200 with Qwen3.5 model.

Key                                                                  | Status  | 
---------------------------------------------------------------------+---------+-
model.language_model.layers.{0...62}.linear_attn.in_proj_qkv.weight  | MISSING | 
model.language_model.layers.{3...63}.self_attn.k_proj.weight         | MISSING | 
model.language_model.layers.{0...62}.linear_attn.dt_bias             | MISSING | 
model.language_model.layers.{3...63}.self_attn.o_proj.weight         | MISSING | 
model.language_model.layers.{0...62}.linear_attn.A_log               | MISSING | 
model.language_model.layers.{0...63}.mlp.up_proj.weight              | MISSING | 
model.language_model.layers.{0...62}.linear_attn.out_proj.weight     | MISSING | 
model.language_model.layers.{0...63}.post_attention_layernorm.weight | MISSING | 
model.language_model.layers.{0...63}.mlp.gate_proj.weight            | MISSING | 
model.language_model.layers.{0...62}.linear_attn.in_proj_b.weight    | MISSING | 
model.language_model.layers.{0...62}.linear_attn.conv1d.weight       | MISSING | 
model.language_model.layers.{3...63}.self_attn.k_norm.weight         | MISSING | 
model.language_model.layers.{3...63}.self_attn.q_proj.weight         | MISSING | 
model.language_model.layers.{0...62}.linear_attn.in_proj_a.weight    | MISSING | 
model.language_model.layers.{0...63}.input_layernorm.weight          | MISSING | 
model.language_model.layers.{0...63}.mlp.down_proj.weight            | MISSING | 
model.language_model.layers.{0...62}.linear_attn.norm.weight         | MISSING | 
model.language_model.layers.{3...63}.self_attn.q_norm.weight         | MISSING | 
model.language_model.layers.{3...63}.self_attn.v_proj.weight         | MISSING | 
model.language_model.layers.{0...62}.linear_attn.in_proj_z.weight    | MISSING | 
model.language_model.embed_tokens.weight                             | MISSING | 
model.language_model.norm.weight                                     | MISSING | 

Edit: transformers 5.5.3, peft 0.18.1, deepspeed 0.18.9, accelerate 1.13.0

[debOliveira]

Also tried the patch at #45402, which seemed like a similar issue, but no success. I'm using my bad fix at the moment.

[zucchini-nlp]

Sorry for long delay, since Cyril is off until tomorrow I tried to reproduce the code and found the root cause. The prev fix is also not completely correct and works because we accidentally end up exploiting another code path

@Cyrilvallez , the regex in qwen3_5_text still matches and replaces sd keys in VLMs. Your PR fixed saving it back, but the loading is broken in DS. This line in loading fn reverts back any conversion, if the converted key is not found in sd (which i believe it not intended). DS load fn, in contrary, doesn't revert any conversion, so it fails to match keys 🙃

transformers/src/transformers/core_model_loading.py

Lines 1219 to 1221 in 4396b1b

	if renamed_key not in meta_model_state_dict and original_key in meta_model_state_dict:
	# Key should probably not have been renamed but we might need the `prefix` to be added.`
	renamed_key, source_pattern = rename_source_key(original_key, [], [], prefix, meta_model_state_dict)

I will let you decide how to fix it when you're back. Personally I believe we shouldn't revert a conversion if not in state dict (I also found that this loophole affected prev CLIP issues, but it's another story)

@debOliveira we'll fix and add it in the next patch, thanks for pinging again :)