Repository: huggingface/transformers
File: src/transformers/integrations/deepspeed.py (Python)

Code before:
    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
    meta_model_state_dict = model.state_dict()

    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, meta_model_state_dict)
            if renamed_key in meta_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 the state dict items to ensure consistent ordering (important for MoE conversions)
    conversion_mapping = {}
    key_rename_cache = {}  # Cache rename results to avoid redundant processing in cleanup loop
    sorted_state_dict = sorted(state_dict.items(), key=lambda kv: dot_natural_key(kv[0]))
    for original_key, tensor in sorted_state_dict:

Review comment (suggestion):
```suggestion
    model_state_dict = model.state_dict()
```

Code after:
    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