Incorrect calculation of strides leading to loss of param data upon tensor parallel use while sliced model loading
Author: kmehant
Pull request: no
Labels: name: Tensor Parallel

name: bug

### System Info

- `transformers` version: 4.50.2
- Python version: 3.12.0
- Huggingface_hub version: 0.29.3
- Safetensors version: 0.5.3
- Accelerate version: 1.5.2
- Accelerate config:    not found
- DeepSpeed version: not installed
- PyTorch version (GPU?): 2.5.0+cu124 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: using TP
- Using GPU in script?: yes
- GPU type: NVIDIA A100-SXM4-80GB

### Who can help?

trainer: @muellerzr @SunMarc
TP: @ArthurZucker

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [x] My own task or dataset (give details below)

### Reproduction

You can load any model (For instance `ibm-granite/granite-3.1-2b-base`) with `tp_plan="auto"`. Make sure the vocab size is not divisible by world size (in this case 5) just for one sharding example, in this case we want to let TP shard lm_head with `colwise_rep`. You can then inspect size of the weight param of lm_head after application and it would be a multiple of world size however loses param data (in this case its downsized from 49152 to 49150). 

### Expected behavior

Loss of data / param parity should hold after application of TP. This is essentially happening due to `get_tensor_shard()` function https://github.com/huggingface/transformers/blob/348f3285c5114159d2ff4933b4b8ae36866d01a7/src/transformers/integrations/tensor_parallel.py#L120 in transformers. Specifically because of flooring of values for instance here - https://github.com/huggingface/transformers/blob/348f3285c5114159d2ff4933b4b8ae36866d01a7/src/transformers/integrations/tensor_parallel.py#L126. Essentially, PyTorch allows for unevening sharding with its placement APIs like Shard.

I am happy to help raise a PR if we agree on an approach. 

Comments:
cc @ArthurZucker 
What kind of solution are you thinking about ? 

Yes! API is missing a few features, and this one is one of them! 
Let's gracefully raise and error if the tensor shape is not divisible! 🤗 happy to review your pr @kmehant 

@ArthurZucker I have added support for this feature instead of raising an error. Please provide your review - https://github.com/huggingface/transformers/pull/37220 Thanks

cc: @SunMarc 

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

@ArthurZucker @SunMarc Requesting review!

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

awaiting for - https://github.com/huggingface/transformers/pull/37220#discussion_r2097381631

Issue and the attached PR are active.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

The issue and attached PR is active and ready for review and merge.

cc: @ArthurZucker @SunMarc