Fix tp (#42368)

ArthurZucker · web-flow · commit 59ed41e6a0c0 · 2025-11-25T10:58:50.000+01:00
* up

* oups you need renamed key for merge moduleoist

* oups there was something I forgot

* update
diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py
@@ -745,7 +745,7 @@ def convert_and_load_state_dict_in_model(
                         mapping.distributed_operation = tp_layer(
                             device_mesh=device_mesh, rank=device_map[""].index, empty_param=empty_param.clone()
                         )
-                    shard_index = len(mapping.collected_tensors)
+                    shard_index = len(mapping.collected_tensors.get(original_key, []))
                     future = spawn_tp_materialize(
                         thread_pool,
                         tensor,

Original file line number	Diff line number	Diff line change
`@@ -745,7 +745,7 @@ def convert_and_load_state_dict_in_model(`
`745`	`745`	`mapping.distributed_operation = tp_layer(`
`746`	`746`	`device_mesh=device_mesh, rank=device_map[""].index, empty_param=empty_param.clone()`
`747`	`747`	`)`
`748`		`- shard_index = len(mapping.collected_tensors)`
	`748`	`+ shard_index = len(mapping.collected_tensors.get(original_key, []))`
`749`	`749`	`future = spawn_tp_materialize(`
`750`	`750`	`thread_pool,`
`751`	`751`	`tensor,`