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

Code before:
    @staticmethod
    def _finalize_model_loading(
        model, load_config: LoadStateDictConfig, loading_info: LoadStateDictInfo
    ) -> LoadStateDictInfo:
        """Perform all post processing operations after having loaded some checkpoints into a model, such as moving
        missing keys from meta device to their expected device, reinitializing missing weights according to proper
        distributions, tying the weights and logging the loading report."""

        # Marks tied weights as `_is_hf_initialized` to avoid initializing them (it's very important for efficiency)
        model.mark_tied_weights_as_initialized()

        model._adjust_missing_keys_with_tied_pointers(loading_info)

        # Move missing (and potentially mismatched) keys and non-persistent buffers back to their expected device from
        # meta device (because they were not moved when loading the weights as they were not in the loaded state dict)
        model._move_missing_keys_from_meta_to_device(
            loading_info.missing_and_mismatched(),
            load_config.device_map,
            load_config.device_mesh,
            load_config.hf_quantizer,
        )
        # Correctly initialize the missing (and potentially mismatched) keys (all parameters without the `_is_hf_initialized` flag)
        model._initialize_missing_keys(load_config.is_quantized)

        # Tie the weights
        model.tie_weights(missing_keys=loading_info.missing_keys)

        # Adjust missing and unexpected keys
        model._adjust_missing_and_unexpected_keys(loading_info)

        log_state_dict_report(
            model=model,
            pretrained_model_name_or_path=load_config.pretrained_model_name_or_path,
            ignore_mismatched_sizes=load_config.ignore_mismatched_sizes,
            loading_info=loading_info,
            logger=logger,
        )

        return loading_info

    def retrieve_modules_from_names(self, names, add_prefix=False, remove_prefix=False):
        module_keys = {".".join(key.split(".")[:-1]) for key in names}

        # torch.nn.ParameterList is a special case where two parameter keywords
        # are appended to the module name, *e.g.* bert.special_embeddings.0
        module_keys = module_keys.union(
            {".".join(key.split(".")[:-2]) for key in names if len(key) > 0 and key[-1].isdigit()}
        )

        retrieved_modules = []
        # retrieve all modules that has at least one missing weight name

Review comment (suggestion):
not sure about this one, it's not BC for models that have `tie_weights` overriden. But we might just say it's supposed to be fixed in remote code

Code after:
            for k in all_pointer:
                k.__exit__(None, None, None)

        return loading_info, disk_offload_index

    @staticmethod
    def _finalize_model_loading(
        model, load_config: LoadStateDictConfig, loading_info: LoadStateDictInfo
    ) -> LoadStateDictInfo:
        """Perform all post processing operations after having loaded some checkpoints into a model, such as moving
        missing keys from meta device to their expected device, reinitializing missing weights according to proper
        distributions, tying the weights and logging the loading report."""
        try:
            # Adjust `all_tied_weights_keys` before marking them as initialized
            model._adjust_tied_keys_with_tied_pointers(loading_info.missing_and_mismatched())

            # Marks tied weights as `_is_hf_initialized` to avoid initializing them (it's very important for efficiency)
            model.mark_tied_weights_as_initialized()

            # Move missing (and potentially mismatched) keys and non-persistent buffers back to their expected device from
            # meta device (because they were not moved when loading the weights as they were not in the loaded state dict)
            model._move_missing_keys_from_meta_to_device(
                loading_info.missing_and_mismatched(),
                load_config.device_map,
                load_config.device_mesh,
                load_config.hf_quantizer,
            )

            # Correctly initialize the missing (and potentially mismatched) keys (all parameters without the `_is_hf_initialized` flag)
            model._initialize_missing_keys(load_config.is_quantized)

            # Tie the weights
            model.tie_weights(missing_keys=loading_info.missing_keys, recompute_mapping=False)

            # Adjust missing and unexpected keys
            model._adjust_missing_and_unexpected_keys(loading_info)
        finally:
            log_state_dict_report(
                model=model,
                pretrained_model_name_or_path=load_config.pretrained_model_name_or_path,
                ignore_mismatched_sizes=load_config.ignore_mismatched_sizes,
                loading_info=loading_info,
                logger=logger,
            )

        return loading_info

    def retrieve_modules_from_names(self, names, add_prefix=False, remove_prefix=False):
        module_keys = {".".join(key.split(".")[:-1]) for key in names}

        # torch.nn.ParameterList is a special case where two parameter keywords