Repository: huggingface/transformers
Pull request: fix TP issue within the device mesh (Tensor Parallel group)


- distributed: @3outeille @ArthurZucker

if you follow readme in https://github.com/huggingface/accelerate/tree/main/examples/torch_native_parallelism to run nd_parallel.py, 
crash like
```
[rank2]: Traceback (most recent call last):
[rank2]:   File "/accelerate/examples/torch_native_parallelism/nd_parallel.py", line 175, in <module>
[rank2]:     train(args)
[rank2]:   File "/accelerate/examples/torch_native_parallelism/nd_parallel.py", line 117, in train
[rank2]:     model = AutoModelForCausalLM.from_pretrained(
[rank2]:             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/transformers/src/transformers/models/auto/auto_factory.py", line 373, in from_pretrained
[rank2]:     return model_class.from_pretrained(
[rank2]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/transformers/src/transformers/modeling_utils.py", line 3991, in from_pretrained
[rank2]:     model, missing_keys, unexpected_keys, mismatched_keys, offload_index, error_msgs = cls._load_pretrained_model(
[rank2]:                                                                                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/transformers/src/transformers/modeling_utils.py", line 4134, in _load_pretrained_model
[rank2]:     convert_and_load_state_dict_in_model(
[rank2]:   File "/transformers/src/transformers/core_model_loading.py", line 1108, in convert_and_load_state_dict_in_model
[rank2]:     realized_value, conversion_errors = mapping.convert(
[rank2]:                                         ^^^^^^^^^^^^^^^^
[rank2]:   File "/transformers/src/transformers/core_model_loading.py", line 581, in convert
[rank2]:     collected_tensors = self.materialize_tensors()
[rank2]:                         ^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/transformers/src/transformers/core_model_loading.py", line 556, in materialize_tensors
[rank2]:     tensors = [future.result() for future in tensors]
[rank2]:                ^^^^^^^^^^^^^^^
[rank2]:   File "/usr/lib/python3.12/concurrent/futures/_base.py", line 449, in result
[rank2]:     return self.__get_result()
[rank2]:            ^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/usr/lib/python3.12/concurrent/futures/_base.py", line 401, in __get_result
[rank2]:     raise self._exception
[rank2]:   File "/usr/lib/python3.12/concurrent/futures/thread.py", line 58, in run
[rank2]:     result = self.fn(*self.args, **self.kwargs)
[rank2]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/transformers/src/transformers/core_model_loading.py", line 721, in _job
[rank2]:     return sharding_method.shard_tensor(tensor, param_casting_dtype=dtype, tensor_idx=tensor_idx)[0]
[rank2]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/transformers/src/transformers/integrations/tensor_parallel.py", line 829, in shard_tensor
[rank2]:     parameter = get_tensor_shard(param, empty_param, device_mesh, rank, -1, tensor_idx=tensor_idx)
[rank2]:                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/transformers/src/transformers/integrations/tensor_parallel.py", line 395, in get_tensor_shard
[rank2]:     raise ValueError(f"Rank {rank} is out of bounds for mesh size {world_size}")
[rank2]: ValueError: Rank 2 is out of bounds for mesh size 2
```

Changed files:
filename: src/transformers/core_model_loading.py
patch: @@ -1064,7 +1064,7 @@ def convert_and_load_state_dict_in_model(
                     if getattr(mapping, "distributed_operation", None) is None:
                         tp_layer = ALL_PARALLEL_STYLES[model.tp_plan[matched_tp_pattern]].__class__
                         mapping.distributed_operation = tp_layer(
-                            device_mesh=device_mesh, rank=device_map[""].index, empty_param=empty_param.clone()
+                            device_mesh=device_mesh, rank=device_mesh.get_local_rank(), empty_param=empty_param.clone()
                         )
                     shard_index = len(mapping.collected_tensors.get(original_key, []))
                     future_or_tensor = spawn_tp_materialize(