Repository: huggingface/transformers
Pull request: [DeepSpeed] properly handle MoE weight conversion

# What does this PR do?

 Fixed the `_apply_weight_conversions_to_state_dict` function to properly handle MoE weight conversions when the renamed key exists in model_state_dict. The bug was that when renamed_key (e.g., `gate_up_proj`) was found in the model's state   dict, the code only counted the match but never added the source tensors (`gate_proj`, `up_proj`) to `conversion_mapping` for fusion.

Fixes #43257 

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker @Cyrilvallez
- vision models: @yonigozlan @molbap
- audio models: @eustlb @ebezzam @vasqu
- multimodal models: @zucchini-nlp
- graph models: @clefourrier

Library:

- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- continuous batching: @remi-or @ArthurZucker @McPatate
- pipelines: @Rocketknight1
- tokenizers: @ArthurZucker and @itazap
- trainer: @SunMarc
- attention: @vasqu @ArthurZucker @CyrilVallez
- model loading (from pretrained, etc): @CyrilVallez
- distributed: @3outeille @ArthurZucker
- CIs: @ydshieh

Integrations:

- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization: @SunMarc @MekkCyber
- kernels: @MekkCyber @drbh
- peft: @BenjaminBossan @githubnemo

Devices/Backends:

- AMD ROCm: @ivarflakstad
- Intel XPU: @IlyasMoutawwakil
- Ascend NPU: @ivarflakstad 

Documentation: @stevhliu

Research projects are not maintained and should be taken as is.

 -->


Changed files:
filename: src/transformers/integrations/deepspeed.py
patch: @@ -348,6 +348,7 @@ def _apply_weight_conversions_to_state_dict(model, state_dict, weight_mapping):
     # Iterate over sorted keys and pop from state_dict to free memory immediately
     conversion_mapping = {}
     key_rename_cache = {}  # Cache rename results to avoid redundant processing
+    new_state_dict = {}  # Initialize here for direct key copies (non-converted keys)
     sorted_keys = sorted(state_dict.keys(), key=lambda k: dot_natural_key(k))
     for original_key in sorted_keys:
         tensor = state_dict.pop(original_key)  # Pop to free memory immediately
@@ -359,20 +360,24 @@ def _apply_weight_conversions_to_state_dict(model, state_dict, weight_mapping):
 
         # Only process if the renamed key is in the model's state dict
         if renamed_key in model_state_dict:
+            # If source_pattern is not None, this key needs WeightConverter (e.g., MoE fusion)
             if source_pattern is not None:
-                new_converter = copy.deepcopy(pattern_to_converter[source_pattern])
+                # Create a fresh converter for this layer to hold its tensors
+                # Share operations list (lightweight, no large data) but get new collected_tensors
+                converter = pattern_to_converter[source_pattern]
+                new_converter = WeightConverter(
+                    source_patterns=converter.source_patterns,
+                    target_patterns=converter.target_patterns,
+                    operations=converter.operations,
+                )
                 mapping = conversion_mapping.setdefault(renamed_key, new_converter)
+                mapping.add_tensor(renamed_key, original_key, source_pattern, tensor)
             else:
-                mapping = conversion_mapping.setdefault(renamed_key, WeightRenaming(original_key, renamed_key))
-                source_pattern = original_key
-
-            # Add the tensor directly (not a Future, since it's already materialized)
-            mapping.add_tensor(renamed_key, original_key, source_pattern, tensor)
+                # No conversion needed - add tensor directly to new_state_dict
+                # (this handles keys like embed_tokens, lm_head, layernorm, attention)
+                new_state_dict[renamed_key] = tensor
 
     # Apply the conversions and build the new state dict
-    new_state_dict = {}
-    # Track which renamed_keys came from WeightConverter (need to skip their originals)
-    converted_renamed_keys = set()
     for renamed_key, mapping in conversion_mapping.items():
         try:
             # Only WeightConverter needs convert(); WeightRenaming is just a simple rename
@@ -386,8 +391,6 @@ def _apply_weight_conversions_to_state_dict(model, state_dict, weight_mapping):
             for target_name, param in realized_value.items():
                 param = param[0] if isinstance(param, list) else param
                 new_state_dict[target_name] = param
-            # Track that this key was converted
-            converted_renamed_keys.add(renamed_key)
             # Free memory by clearing source tensors
             if hasattr(mapping, "source_tensors"):
                 mapping.source_tensors = {}

filename: tests/deepspeed/test_deepspeed.py
patch: @@ -476,6 +476,9 @@ def test_init_zero3_moe_weight_conversion(self):
                         loaded_model = Qwen3MoeModel.from_pretrained(old_checkpoint_dir)
 
             self.assertIn("Detected DeepSpeed ZeRO-3", cl.out)
+
+            # Without weight conversion, gate_up_proj and down_proj would be MISSING
+            # This regex fails the test if expert fusion weights are missing from checkpoint
             self.assertNotRegex(cl.out, r"mlp\.experts\.(gate_up_proj|down_proj)\s*\|\s*MISSING")
 
             # Verify the model structure is correct (fused experts in v5 format)
@@ -488,6 +491,14 @@ def test_init_zero3_moe_weight_conversion(self):
                     expert_params_to_check.append((name, param))
                 self.assertNotRegex(name, r"mlp\.experts\.\d+\.(gate_proj|up_proj|down_proj)\.weight")
 
+            # Without the fix, expert_params_to_check would be empty (all MISSING)
+            self.assertGreater(
+                len(expert_params_to_check),
+                0,
+                "No expert weights found - weight conversion failed! "
+                "Expected fused gate_up_proj and down_proj but found none.",
+            )
+
             with deepspeed.zero.GatheredParameters([param for _, param in expert_params_to_check], modifier_rank=0):
                 for name, param in expert_params_to_check:
                     if "mlp.experts.gate_up_proj" in name: