From 12eb8113bb1f06b0f89b7cc651baea4b69a38051 Mon Sep 17 00:00:00 2001 From: Cyril Vallez Date: Wed, 20 May 2026 18:03:50 +0900 Subject: [PATCH] improve --- src/transformers/core_model_loading.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index e15c2314ac31..cd70209e0cb1 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -161,7 +161,10 @@ def convert( for source_pattern in source_patterns: if source_pattern not in input_dict: continue - tensors = input_dict[source_pattern] + # Immediately free the input_dict, so that we do not keep many copies simultaneously - otherwise we have to + # wait for this function to return to be able to clean-up, which will not get garbage collected as fast as if + # everything is freed right now + tensors = input_dict.pop(source_pattern) if isinstance(tensors, list): all_tensors.extend(tensors) else: @@ -197,15 +200,20 @@ def convert( target_patterns: list[str], **kwargs, ) -> dict[str, torch.Tensor]: + input_size = len(input_dict) merged: dict[str, torch.Tensor] = {} - for source_pattern, tensors in input_dict.items(): - target_pattern = self.get_target_pattern(input_dict, source_pattern, target_patterns) + for source_pattern in list(input_dict.keys()): + # Immediately free the input dict, so that we do not keep many copies simultaneously - otherwise we have to + # wait for this function to return to be able to clean-up, and if the size of the input_dict is larger than 1 + # (such as the MoEs' gate_proj/up_proj merging), we are wasting quite some memory + tensors = input_dict.pop(source_pattern) + target_pattern = self.get_target_pattern(input_size, source_pattern, target_patterns) merged[target_pattern] = torch.stack(tensors, dim=self.dim) return merged - def get_target_pattern(self, input_dict: dict, source_pattern: str, target_patterns: list[str]) -> str: + def get_target_pattern(self, input_size: int, source_pattern: str, target_patterns: list[str]) -> str: # Here it's a single operation, so we use the target - if len(input_dict) == 1: + if input_size == 1: if len(target_patterns) == 1: return target_patterns[0] else: