Loading optimization (#42239)

Cyrilvallez · web-flow · commit 0e74a71c03b7 · 2025-11-17T14:28:57.000+01:00
* remove loop over modules

* no need for set casting
diff --git 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 = 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:

Original file line number	Diff line number	Diff line change
`@@ -4709,7 +4709,7 @@ def caching_allocator_warmup(model: PreTrainedModel, expanded_device_map: dict,`
`4709`	`4709`	`else None`
`4710`	`4710`	`)`
`4711`	`4711`	`total_byte_count = defaultdict(lambda: 0)`
`4712`		`- tied_param_names = _get_tied_weight_keys(model)`
	`4712`	`+ tied_param_names = model.all_tied_weights_keys.keys()`
`4713`	`4713`	`for param_name, device in accelerator_device_map.items():`
`4714`	`4714`	`# Skip if the parameter has already been accounted for (tied weights)`
`4715`	`4715`	`if param_name in tied_param_names:`