From 46f442973831b8e054baefe0fb33b565ff6e16a5 Mon Sep 17 00:00:00 2001 From: Cyril Vallez Date: Mon, 17 Nov 2025 13:29:52 +0100 Subject: [PATCH 1/2] remove loop over modules --- src/transformers/modeling_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 965084d0c24a..a2111bf2743e 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -4709,7 +4709,7 @@ def caching_allocator_warmup(model: PreTrainedModel, expanded_device_map: dict, else None ) total_byte_count = defaultdict(lambda: 0) - tied_param_names = _get_tied_weight_keys(model) + tied_param_names = set(model.all_tied_weights_keys.keys()) for param_name, device in accelerator_device_map.items(): # Skip if the parameter has already been accounted for (tied weights) if param_name in tied_param_names: From fc5970bb4f1c2455135a143d39d82774e305dfdd Mon Sep 17 00:00:00 2001 From: Cyril Vallez Date: Mon, 17 Nov 2025 13:31:49 +0100 Subject: [PATCH 2/2] no need for set casting --- src/transformers/modeling_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index a2111bf2743e..3df4f6d47b5a 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -4709,7 +4709,7 @@ def caching_allocator_warmup(model: PreTrainedModel, expanded_device_map: dict, else None ) total_byte_count = defaultdict(lambda: 0) - tied_param_names = set(model.all_tied_weights_keys.keys()) + tied_param_names = model.all_tied_weights_keys.keys() for param_name, device in accelerator_device_map.items(): # Skip if the parameter has already been accounted for (tied weights) if param_name in tied_param_names: