Repository: huggingface/transformers
Pull request: [DeepSpeed] add weight_mapping to _load_state_dict_into_zero3_model

# What does this PR do?

This pull request introduces enhancements to the DeepSpeed integration for model loading, specifically improving how weight conversions (such as renaming and merging/splitting of weights) are handled when loading state dictionaries into models using DeepSpeed ZeRO-3. The changes add a new utility for applying weight conversion logic before loading, and update the model loading pipeline to use this utility when needed.

Key changes:

### Weight Conversion Utilities

* Added a new function `_apply_weight_conversions_to_state_dict` in `deepspeed.py` to handle weight renaming and merging/splitting operations on state dicts before loading them into a model. This function supports both simple renaming and more complex conversions using `WeightConverter` and `WeightRenaming` objects.

### DeepSpeed Model Loading Logic

* Updated `_load_state_dict_into_zero3_model` to accept an optional `weight_mapping` argument. If provided, it uses the new `_apply_weight_conversions_to_state_dict` function to preprocess the state dict before loading. It also stores the applied weight conversions on the model for later reference.
* Modified the call to `_load_state_dict_into_zero3_model` in `_load_pretrained_model` (in `modeling_utils.py`) to pass the `weight_mapping` argument, ensuring that weight conversions are applied during DeepSpeed ZeRO-3 model loading.

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: @@ -290,19 +290,160 @@ def deepspeed_config():
         return None
 
 
-def _load_state_dict_into_zero3_model(model_to_load, state_dict):
+def _apply_weight_conversions_to_state_dict(model, state_dict, weight_mapping):
+    """
+    Apply weight conversions (renaming and merging/splitting operations) to a state dict.
+    This is a simplified version that handles the conversion without loading into the model.
+    """
+    # Check for Tensor Parallelism - weight conversions are not tested with TP
+    # TP uses ReplaceWithTensorSlicing which may conflict with our weight conversions
+    ds_config = deepspeed_config()
+    if ds_config is not None:
+        # Check training config (tensor_parallel.autotp_size)
+        tp_size = ds_config.get("tensor_parallel", {}).get("autotp_size", 1)
+        # Check inference config (inference.tensor_parallel.tp_size)
+        inference_config = ds_config.get("inference", {})
+        if isinstance(inference_config, dict):
+            tp_size = max(tp_size, inference_config.get("tensor_parallel", {}).get("tp_size", 1))
+        if tp_size > 1:
+            raise NotImplementedError(
+                "Weight conversions (e.g., MoE expert fusion) with DeepSpeed Tensor Parallelism "
+                "are not yet implemented but support is coming soon. Please disable tensor_parallel "
+                "in your DeepSpeed config or convert your checkpoint to the expected format first."
+            )
+
+    from ..core_model_loading import WeightConverter, WeightRenaming, dot_natural_key, rename_source_key
+
+    # Preserve metadata from the original state dict
+    metadata = getattr(state_dict, "_metadata", None)
+
+    prefix = model.base_model_prefix
+
+    # Build a meta state dict for matching - only keys/shapes, no actual tensor data
+    # This minimizes memory since we don't duplicate the model's parameters
+    model_state_dict = {}
+    for key, param in model.state_dict().items():
+        model_state_dict[key] = torch.empty(param.shape, dtype=param.dtype, device="meta")
+
+    renamings = [entry for entry in weight_mapping if isinstance(entry, WeightRenaming)]
+    converters = [entry for entry in weight_mapping if isinstance(entry, WeightConverter)]
+
+    # Fast path: if we only have simple renamings and no converters, we can skip the expensive collection logic
+    if len(converters) == 0:
+        new_state_dict = {}
+        for original_key, tensor in state_dict.items():
+            renamed_key, _ = rename_source_key(original_key, renamings, [], prefix, model_state_dict)
+            if renamed_key in model_state_dict:
+                new_state_dict[renamed_key] = tensor
+        # Attach metadata to the new state dict
+        if metadata is not None:
+            new_state_dict._metadata = metadata
+        return new_state_dict
+
+    # Full path: we have WeightConverter operations that require tensor fusion/splitting
+    pattern_to_converter = {k: converter for converter in converters for k in converter.source_patterns}
+
+    # Build a mapping of what needs to be converted
+    # 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
+    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
+        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:
+                new_converter = copy.deepcopy(pattern_to_converter[source_pattern])
+                mapping = conversion_mapping.setdefault(renamed_key, new_converter)
+            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)
+
+    # 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
+            if not isinstance(mapping, WeightConverter):
+                continue
+            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
+            # 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 = {}
+        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
+
+    return new_state_dict
+
+
+def _load_state_dict_into_zero3_model(model_to_load, state_dict, load_config=None):
     """
     Loads state dict into a model specifically for Zero3, since DeepSpeed does not support the `transformers`
     tensor parallelism API.
 
     Nearly identical code to PyTorch's `_load_from_state_dict`
+
+    Args:
+        model_to_load: The model to load weights into
+        state_dict: The state dict containing the weights
+        load_config: Optional LoadStateDictConfig containing weight_mapping and other loading options
     """
     # copy state_dict so `_load_state_dict_into_zero3_model` can modify it
     metadata = getattr(state_dict, "_metadata", None)
     state_dict = state_dict.copy()
     if metadata is not None:
         state_dict._metadata = metadata
 
+    # Extract weight_mapping from load_config if provided
+    weight_mapping = None
+    if load_config is not None:
+        weight_mapping = getattr(load_config, "weight_mapping", None)
+
+    # Apply weight conversions if provided
+    if weight_mapping is not None and len(weight_mapping) > 0:
+        state_dict = _apply_weight_conversions_to_state_dict(model_to_load, state_dict, weight_mapping)
+        # Keep the current weight conversion mapping for later saving (in case it was coming directly from the user)
+        model_to_load._weight_conversions = weight_mapping
+
     error_msgs = []
     meta_model_state_dict = model_to_load.state_dict()
     missing_keys = set(meta_model_state_dict.keys())

filename: src/transformers/modeling_utils.py
patch: @@ -4205,7 +4205,7 @@ def _load_pretrained_model(
                         load_state_dict(ckpt_file, map_location="cpu", weights_only=load_config.weights_only)
                     )
                 state_dict = merged_state_dict
-            error_msgs, missing_keys = _load_state_dict_into_zero3_model(model, state_dict)
+            error_msgs, missing_keys = _load_state_dict_into_zero3_model(model, state_dict, load_config)
             # This is not true but for now we assume only best-case scenario with deepspeed, i.e. perfectly matching checkpoints
             unexpected_keys, mismatched_keys, conversion_errors = set(), set(), set()
         else:

filename: tests/deepspeed/test_deepspeed.py
patch: @@ -55,6 +55,8 @@
 
 
 if is_torch_available():
+    import os
+
     import torch
     import torch.nn as nn
 
@@ -387,6 +389,117 @@ def test_arange_bf16(self):
         # check that we get the same results either with torch or deepspeed
         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
+        import tempfile
+
+        from transformers import Qwen3MoeConfig, Qwen3MoeModel
+
+        tiny_config = Qwen3MoeConfig(
+            vocab_size=99,
+            hidden_size=32,
+            intermediate_size=32,
+            moe_intermediate_size=16,
+            num_hidden_layers=2,
+            num_attention_heads=2,
+            num_key_value_heads=2,
+            num_experts=8,
+            num_experts_per_tok=2,
+        )
+
+        ds_config = {
+            "train_batch_size": 1,
+            "zero_optimization": {
+                "stage": 3,
+            },
+        }
+
+        dschf = HfDeepSpeedConfig(ds_config)
+
+        self.assertTrue(dschf.is_zero3())
+        self.assertTrue(is_deepspeed_zero3_enabled())
+
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            with LoggingLevel(logging.INFO):
+                with mockenv_context(**self.dist_env_1_gpu):
+                    model = Qwen3MoeModel(tiny_config)
+                    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)
+
+            self.assertIn("Detected DeepSpeed ZeRO-3", cl.out)
+            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")
+
+            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]}")
+
 
 class TrainerIntegrationDeepSpeedWithCustomConfig(TestCasePlus):
     def setUp(self):