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

Code before:
    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:
        # 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, meta_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 meta_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 = {}
    for first_param_name, mapping in conversion_mapping.items():
        try:
            realized_value, _ = mapping.convert(
                first_param_name,
                model=model,
                config=model.config,

Review comment (suggestion):
```suggestion
            original_key, renamings, converters, prefix, model_state_dict
```

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