feat: add support for tensor parallel training workflow with accelerate
Author: kmehant
Pull request: yes

# What does this PR do?

~1. Add `apply_tensor_parallel` API to apply TP plan to Llama and Granite models~ Already merged in to HF/transformers.
2. Introduce `tp_size` user facing argument to be further consumed by accelerate (see https://github.com/huggingface/accelerate/pull/3173)
3. Allows for e2e TP training.

Please review in conjunction with https://github.com/huggingface/accelerate/pull/3173

Fixes https://github.com/huggingface/transformers/issues/32470

# Results

See significant improvement in both memory and throughput compared against single gpu training, and FSDP across different settings (checkpointing on/off) and context lengths.

**Note:** Please be aware that the effective TPS for FSDP would be multiplicative of the parallel factor (number of GPUs/devices engaged in distributed training) whereas that is not the case with TP. Therefore, when effective throughput is considered we can find FSDP is better than TP in terms of throughput. However, that **may** be compensated by increasing the batch size utilizing the memory gains etc. 

Done on two models
1. ibm-granite/granite-8b-code-base-128k 
4. codellama/CodeLlama-7b-hf

Tables below show the max cuda memory and throughput for various configurations showing the potential of TP contributed in this PR. There is gains in both memory and throughput.

Model | Method | # of GPUs | Context Length | Batch Size | Grad Checkpointing | Cuda Max Mem (GiB) | Tokens/Sec/GPU
-- | -- | -- | -- | -- | -- | -- | --
ibm-granite/granite-8b-code-base-128k | Single GPU non-distributed | 1 | 8192 | 1 | FALSE | OOM | NA
ibm-granite/granite-8b-code-base-128k | FSDP | 4 | 8192 | 1 | FALSE | OOM | NA
ibm-granite/granite-8b-code-base-128k | TP (This PR) | 4 | 8192 | 1 | FALSE | 52.4 | 7675.4

Model | Method | # of GPUs | Context Length | Batch Size | Grad Checkpointing | Cuda Max Mem (GiB) | Tokens/Sec/GPU
-- | -- | -- | -- | -- | -- | -- | --
ibm-granite/granite-8b-code-base-128k | Single GPU non-distributed | 1 | 8192 | 1 | TRUE | OOM | NA
ibm-granite/granite-8b-code-base-128k | FSDP | 4 | 8192 | 1 | TRUE | 29.975586 | 2256.896
ibm-granite/granite-8b-code-base-128k | TP (This PR) | 4 | 8192 | 1 | TRUE | 26.5 | 5935.5

Model | Method | # of GPUs | Context Length | Batch Size | Grad Checkpointing | Cuda Max Mem (GiB) | Tokens/Sec/GPU
-- | -- | -- | -- | -- | -- | -- | --
ibm-granite/granite-8b-code-base-128k | Single GPU non-distributed | 1 | 16384 | 1 | FALSE | OOM | NA
ibm-granite/granite-8b-code-base-128k | FSDP | 4 | 16384 | 1 | FALSE | OOM | NA
ibm-granite/granite-8b-code-base-128k | TP (This PR) | 4 | 16384 | 1 | FALSE | OOM | NA

Model | Method | # of GPUs | Context Length | Batch Size | Grad Checkpointing | Cuda Max Mem (GiB) | Tokens/Sec/GPU
-- | -- | -- | -- | -- | -- | -- | --
ibm-granite/granite-8b-code-base-128k | Single GPU non-distributed | 1 | 16384 | 1 | TRUE | OOM | NA
ibm-granite/granite-8b-code-base-128k | FSDP | 4 | 16384 | 1 | TRUE | 36.8 | 2084.864
ibm-granite/granite-8b-code-base-128k | TP (This PR) | 4 | 16384 | 1 | TRUE | 33.5 | 5692.5


Model | Method | # of GPUs | Context Length | Batch Size | Grad Checkpointing | Cuda Max Mem (GiB) | Tokens/Sec/GPU
-- | -- | -- | -- | -- | -- | -- | --
codellama/CodeLlama-7b-hf | Single GPU non-distributed | 1 | 8192 | 1 | FALSE | OOM | NA
codellama/CodeLlama-7b-hf | FSDP | 4 | 8192 | 1 | FALSE | 70.7 | 3560
codellama/CodeLlama-7b-hf | TP (This PR) | 4 | 8192 | 1 | FALSE | 42.8 | 9216

Model | Method | # of GPUs | Context Length | Batch Size | Grad Checkpointing | Cuda Max Mem (GiB) | Tokens/Sec/GPU
-- | -- | -- | -- | -- | -- | -- | --
codellama/CodeLlama-7b-hf | Single GPU non-distributed | 1 | 8192 | 1 | TRUE | 75.3 | 2849
codellama/CodeLlama-7b-hf | FSDP | 4 | 8192 | 1 | TRUE | 26.4 | 5957
codellama/CodeLlama-7b-hf | TP (This PR) | 4 | 8192 | 1 | TRUE | 21.4 | 7125

Model | Method | # of GPUs | Context Length | Batch Size | Grad Checkpointing | Cuda Max Mem (GiB) | Tokens/Sec/GPU
-- | -- | -- | -- | -- | -- | -- | --
codellama/CodeLlama-7b-hf | Single GPU non-distributed | 1 | 16384 | 1 | FALSE | OOM | NA
codellama/CodeLlama-7b-hf | FSDP | 4 | 16384 | 1 | FALSE | OOM | NA
codellama/CodeLlama-7b-hf | TP (This PR) | 4 | 16384 | 1 | FALSE | OOM | NA

Model | Method | # of GPUs | Context Length | Batch Size | Grad Checkpointing | Cuda Max Mem (GiB) | Tokens/Sec/GPU
-- | -- | -- | -- | -- | -- | -- | --
codellama/CodeLlama-7b-hf | Single GPU non-distributed | 1 | 16384 | 1 | TRUE | 75.3 | 2599
codellama/CodeLlama-7b-hf | FSDP | 4 | 16384 | 1 | TRUE | 30.1 | 2433
codellama/CodeLlama-7b-hf | TP (This PR) | 4 | 16384 | 1 | TRUE | 26.6 | 6873


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

- trainer: @muellerzr and @SunMarc 
- @SeungyounShin (happy to include as co author - https://github.com/huggingface/transformers/pull/32597)

I have cycles to bring in more improvements over this PR to bring in Pytorch TP support to HF. Looking forward. Thank you

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)


Comments:
Such timing! I have similar thought [here](https://github.com/huggingface/transformers/pull/34184). Shall we collaborate? 

@kwen2501 Absolutely, please let me know, how you want to take this forward. Thank you.

@ArthurZucker @muellerzr since accelerate PR (https://github.com/huggingface/accelerate/pull/3173) is merged. Requesting review and merge of this PR which would allow for complete e2e training workflow using tensor parallelism. Thank you. 

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_34194). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@SunMarc Have added some documentation. I would also add a page of documentation in accelerate and then reference it back here in a separate PR. For now I kept the documentation in the PR self contained for merge and use. Thanks. 

cc: @muellerzr 

@SunMarc failing tests look unrelated to this PR and seem to happen after I rebased with main branch. Thank you. 

@SunMarc @muellerzr I have a follow up PR as well - https://github.com/huggingface/transformers/pull/36132 FYA. Thank you.

@SunMarc for some reason, even for the recently merged commits are failing for these two CI tests.

Very much welcome otherwise!!!


@ArthurZucker Have added and updated docs wherever needed for TP, requesting your review. Thank you.

@ArthurZucker failing tests do not seem to be related to this PR

@ArthurZucker @muellerzr @SunMarc Can we go ahead with merge?

Hi @kmehant thank you very much for the much needed PR! Just wanna get clarified, aside from setting the `tp_size` which is passed into the transformers trainer, you also need to set  `tp_plan="auto"` in the `automodel.from_pretrained(...)` right?

@bursteratom setting `tp_plan="auto"` is not needed since TP is applied to the model internally by accelerate when you are using through transformers with `tp_size`.

@kmehant Thank you for your response! I am asking cuz I ran into `RuntimeError: aten._foreach_norm.Scalar: got mixed torch.Tensor and DTensor, need to convert all torch.Tensor to DTe
nsor before calling distributed operators!` when attempting to do TP with accelerate. I only set `tp_size=2` in the trainer kwarg and did not set `tp_plan="auto"` at all. I am wondering if there are some additional steps that I missed? (testing with `codellama/CodeLlama-7b-hf` btw)

@bursteratom Right, apologies for not covering limitations in this PR. Essentially this happening due to gradient clipping which is not supported yet (PR is in place to fix this - https://github.com/huggingface/transformers/pull/36132). So a workaround for this would be to set `max_grad_norm to -1` until the above PR falls in place. 

cc: @kwen2501 @muellerzr @SunMarc 

@kmehant Thank you so much for the tip and your amazing work! Looking forward to your fix getting merged!

@kmehant apologies for the repeated ping. I was able to train with tensor parallel successfully in accelerate, but ran into `RuntimeError: Attempted to access the data pointer on an invalid python storage.` when attempting to save the trained model. Any chance you know why this is happening?
 
![image](https://github.com/user-attachments/assets/50982748-f83b-4748-94bd-a1f11e62ff30)