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

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


def _load_state_dict_into_zero3_model(model_to_load, state_dict, weight_mapping=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
        weight_mapping: Optional list of WeightConverter/WeightRenaming operations to apply
    """
    # 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

    # 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())

    prefix_model = getattr(model_to_load, "base_model_prefix", None)
    # take care of the case where in the checkpoint we don't have the prefix
    state_dict = {
        (f"{prefix_model}.{k}" if meta_model_state_dict.get(f"{prefix_model}.{k}") is not None else k): v
        for k, v in state_dict.items()
    }

    # PyTorch's `_load_from_state_dict` does not copy parameters in a module's descendants
    # so we need to apply the function recursively.
    def load(module: nn.Module, state_dict, prefix="", assign_to_params_buffers=False):
        local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
        local_metadata["assign_to_params_buffers"] = assign_to_params_buffers

        args = (state_dict, prefix, local_metadata, True, [], [], error_msgs)

Review comment (suggestion):
we need to delete the old value otherwise we are gonna explose memory consumption. We don't have to in the other path because we read only once

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