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

Code before:
        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, this key needs WeightConverter (e.g., MoE fusion)
            if source_pattern is not None:
                # Create a fresh converter for this layer to hold its tensors
                # Share operations list (lightweight, no large data) but get new collected_tensors
                converter = pattern_to_converter[source_pattern]
                new_converter = WeightConverter(
                    source_patterns=converter.source_patterns,
                    target_patterns=converter.target_patterns,
                    operations=converter.operations,
                )
                mapping = conversion_mapping.setdefault(renamed_key, new_converter)
                mapping.add_tensor(renamed_key, original_key, source_pattern, tensor)
            else:
                # No conversion needed - add tensor directly to new_state_dict
                # (this handles keys like embed_tokens, lm_head, layernorm, attention)
                new_state_dict[renamed_key] = tensor
        else:
            if source_pattern is not None:
                # Create a fresh converter for this layer to hold its tensors
                # Share operations list (lightweight, no large data) but get new collected_tensors
                converter = pattern_to_converter[source_pattern]
                new_converter = WeightConverter(
                    source_patterns=converter.source_patterns,
                    target_patterns=converter.target_patterns,
                    operations=converter.operations,
                )
                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
    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,

Review comment (suggestion):
I am actually not sure we should do the else here. If the renamed key is not in model_state_dict then I don't get how it should be loaded. The "contract" is that the model is pre-processed first to have all the final layers (quantization / peft etc) and then we just have to check the keys of tthe state dict to be good!

Code after:
        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, this key needs WeightConverter (e.g., MoE fusion)
            if source_pattern is not None:
                # Create a fresh converter for this layer to hold its tensors
                # Share operations list (lightweight, no large data) but get new collected_tensors
                converter = pattern_to_converter[source_pattern]
                new_converter = WeightConverter(
                    source_patterns=converter.source_patterns,
                    target_patterns=converter.target_patterns,
                    operations=converter.operations,
                )
                mapping = conversion_mapping.setdefault(renamed_key, new_converter)
                mapping.add_tensor(renamed_key, original_key, source_pattern, tensor)
            else:
                # No conversion needed - add tensor directly to new_state_dict
                # (this handles keys like embed_tokens, lm_head, layernorm, attention)
                new_state_dict[renamed_key] = tensor

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