Repository: huggingface/transformers
Pull request: [Deespeed] fix WeightConverter.convert() use

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

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: @@ -347,17 +347,12 @@ def _apply_weight_conversions_to_state_dict(model, state_dict, weight_mapping):
     # Sort keys to ensure consistent ordering (important for MoE conversions)
     # 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)
+    new_state_dict = {}
     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
-        # Rename the key according to all renaming pattern and optional weight converter patterns
+        tensor = state_dict.pop(original_key)
         renamed_key, source_pattern = rename_source_key(original_key, renamings, converters, prefix, model_state_dict)
 
-        # Cache the rename result for use in the cleanup loop
-        key_rename_cache[original_key] = renamed_key
-
         # 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)
@@ -380,37 +375,21 @@ def _apply_weight_conversions_to_state_dict(model, state_dict, weight_mapping):
     # Apply the conversions and build the new state dict
     for renamed_key, mapping in conversion_mapping.items():
         try:
-            # Only WeightConverter needs convert(); WeightRenaming is just a simple rename
-            if not isinstance(mapping, WeightConverter):
-                continue
-            realized_value, _ = mapping.convert(
+            realized_value = mapping.convert(
                 renamed_key,
                 model=model,
                 config=model.config,
             )
             for target_name, param in realized_value.items():
                 param = param[0] if isinstance(param, list) else param
                 new_state_dict[target_name] = param
-            # Free memory by clearing source tensors
-            if hasattr(mapping, "source_tensors"):
-                mapping.source_tensors = {}
         except Exception as e:
             raise RuntimeError(
                 f"Failed to apply weight conversion for '{renamed_key}'. "
                 f"This likely means the checkpoint format is incompatible with the current model version. "
                 f"Error: {e}"
             ) from e
 
-    # Add any keys that didn't need conversion (use cached rename results)
-    # At this point, state_dict only contains unconverted keys (others were popped)
-    for key in list(state_dict.keys()):
-        renamed_key = key_rename_cache.get(key)
-        if renamed_key is None:
-            # Key wasn't in our cache, compute rename
-            renamed_key, _ = rename_source_key(key, renamings, [], prefix, model_state_dict)
-        if renamed_key not in new_state_dict and renamed_key in model_state_dict:
-            new_state_dict[renamed_key] = state_dict.pop(key)
-
     # Attach metadata to the new state dict
     if metadata is not None:
         new_state_dict._metadata = metadata

filename: tests/deepspeed/test_deepspeed.py
patch: @@ -390,11 +390,12 @@ def test_arange_bf16(self):
         torch.testing.assert_close(good_torch_sin_cos, good_deepspeed_sin_cos.cpu())
 
     def test_init_zero3_moe_weight_conversion(self):
-        # Test that weight conversions (MoE expert fusion) work correctly with DeepSpeed Zero3
-        # This tests the fix for the issue where DeepSpeed Zero3 loading was bypassing weight conversions
+        # test that weight conversions (MoE expert fusion) work correctly under zero3
         import tempfile
 
-        from transformers import Qwen3MoeConfig, Qwen3MoeModel
+        import deepspeed
+
+        from transformers import Qwen3MoeConfig, Qwen3MoeForCausalLM
 
         tiny_config = Qwen3MoeConfig(
             vocab_size=99,
@@ -423,93 +424,32 @@ def test_init_zero3_moe_weight_conversion(self):
         with tempfile.TemporaryDirectory() as tmpdirname:
             with LoggingLevel(logging.INFO):
                 with mockenv_context(**self.dist_env_1_gpu):
-                    model = Qwen3MoeModel(tiny_config)
+                    model = Qwen3MoeForCausalLM(tiny_config)
+                    reference_weights = {k: v.clone() for k, v in model.state_dict().items()}
                     model.save_pretrained(tmpdirname)
 
-            # Manually create an "old" checkpoint format with separate expert weights
-            # to simulate loading a checkpoint saved before the v5 refactor
-            old_checkpoint_dir = f"{tmpdirname}_old"
-            import os
-            import shutil
-
-            os.makedirs(old_checkpoint_dir, exist_ok=True)
-            shutil.copy(f"{tmpdirname}/config.json", f"{old_checkpoint_dir}/config.json")
-
-            from safetensors.torch import load_file, save_file
-
-            new_state_dict = load_file(f"{tmpdirname}/model.safetensors")
-            old_state_dict = {}
-
-            for key, tensor in new_state_dict.items():
-                if "mlp.experts.gate_up_proj" in key:
-                    layer_prefix = key.replace(".mlp.experts.gate_up_proj", "")
-                    num_experts = tensor.shape[0]
-                    intermediate_size = tensor.shape[1] // 2
-
-                    for expert_idx in range(num_experts):
-                        expert_tensor = tensor[expert_idx]
-                        gate_tensor = expert_tensor[:intermediate_size, :]
-                        up_tensor = expert_tensor[intermediate_size:, :]
-
-                        old_state_dict[f"{layer_prefix}.mlp.experts.{expert_idx}.gate_proj.weight"] = gate_tensor
-                        old_state_dict[f"{layer_prefix}.mlp.experts.{expert_idx}.up_proj.weight"] = up_tensor
-                elif (
-                    "mlp.experts.down_proj" in key
-                    and key[key.rfind(".mlp.experts.down_proj") :] == ".mlp.experts.down_proj"
-                ):
-                    layer_prefix = key.replace(".mlp.experts.down_proj", "")
-                    num_experts = tensor.shape[0]
-
-                    for expert_idx in range(num_experts):
-                        expert_tensor = tensor[expert_idx]
-                        old_state_dict[f"{layer_prefix}.mlp.experts.{expert_idx}.down_proj.weight"] = expert_tensor
-                else:
-                    old_state_dict[key] = tensor
-
-            save_file(old_state_dict, f"{old_checkpoint_dir}/model.safetensors")
-
-            # Load the old checkpoint with DeepSpeed Zero3 and verify weight conversions are applied
             with LoggingLevel(logging.INFO):
                 with mockenv_context(**self.dist_env_1_gpu):
                     logger = logging.get_logger("transformers.modeling_utils")
                     with CaptureLogger(logger) as cl:
-                        loaded_model = Qwen3MoeModel.from_pretrained(old_checkpoint_dir)
-
+                        loaded_model, loading_info = Qwen3MoeForCausalLM.from_pretrained(
+                            tmpdirname, output_loading_info=True
+                        )
             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)
-            # DeepSpeed Zero3 partitions parameters, so we need to gather them to check shapes
-            import deepspeed
-
-            expert_params_to_check = []
-            for name, param in loaded_model.named_parameters():
-                if "mlp.experts.gate_up_proj" in name or "mlp.experts.down_proj" in name:
-                    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.",
+            self.assertEqual(len(loading_info["missing_keys"]), 0, f"Missing keys: {loading_info['missing_keys']}")
+            self.assertEqual(
+                len(loading_info["unexpected_keys"]), 0, f"Unexpected keys: {loading_info['unexpected_keys']}"
             )
 
-            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:
-                        self.assertEqual(len(param.shape), 3, f"gate_up_proj should be 3D, got {param.shape}")
-                        self.assertEqual(param.shape[0], 8, f"Should have 8 experts, got {param.shape[0]}")
-                    elif (
-                        "mlp.experts.down_proj" in name
-                        and name[name.rfind(".mlp.experts.down_proj") :] == ".mlp.experts.down_proj"
-                    ):
-                        self.assertEqual(len(param.shape), 3, f"down_proj should be 3D, got {param.shape}")
-                        self.assertEqual(param.shape[0], 8, f"Should have 8 experts, got {param.shape[0]}")
+            # gather all params and verify they match the original weights exactly
+            all_params = list(loaded_model.named_parameters())
+            with deepspeed.zero.GatheredParameters([p for _, p in all_params], modifier_rank=0):
+                for name, param in all_params:
+                    torch.testing.assert_close(
+                        param.data.cpu(),
+                        reference_weights[name].cpu(),
+                        msg=f"Parameter '{name}' doesn't match reference weights",
+                    )
 
     def test_init_zero3_variance_scaling(self):
         """