Simplify Tensor Parallel implementation with PyTorch TP
Author: kwen2501
Pull request: yes

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

This PR uses the `torch.distributed.tensor.parallel` subpackage to implement Tensor Parallel for Llama (as an example).

The motivation is multi-fold:

1. to make modeling code simple as single-worker case:
all manual TP implementations under `if self.config.pretraining_tp > 1` can be removed.

3. to make tensor parallelism easily accessible by users:
added a `model.tensor_parallel(device_mesh)` method that allows users to turn a single-proc model into a parallel model. !- Please guide me to a right place to put this function/method if `PreTrainedModel` is not a preferred place. -!

### Note: 
This is just a demo. The removal of `if self.config.pretraining_tp > 1` may break workflows elsewhere, but it is just intended to calculate how much code can be saved, and hopefully it would be possible to direct it to the mechanism in this PR.

### User script:
```
import os
import torch
import torch.distributed as dist

from transformers import AutoTokenizer, AutoModelForCausalLM

model_id = "meta-llama/Meta-Llama-3-8B-Instruct"

def main(rank, world_size) -> None:
    device = torch.device(f"cuda:{rank}")
    dist.init_process_group("nccl", device_id=device)

    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        tp_plan="auto",
    )

    tokenizer = AutoTokenizer.from_pretrained(model_id)
    prompt = "Can I help"

    inputs = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
    outputs = model(inputs)

    next_token_logits = outputs[0][:, -1, :]
    next_token = torch.argmax(next_token_logits, dim=-1)
    response = tokenizer.batch_decode(next_token)
    print(f"{rank=}: {response}")


if __name__ == "__main__":
    main(int(os.environ["RANK"]), int(os.environ["WORLD_SIZE"]))

```

Launch command:
`torchrun --standalone --nproc-per-node 4 tp_hf.py`

Output:
```
Can I help
rank=3: [' you']
rank=1: [' you']
rank=2: [' you']
rank=0: [' you']
```

### Test
```
CUDA_VISIBLE_DEVICES=0,1 pytest -sv tests/tp/test_tp.py
```

### Doc
`perf_infer_gpu_multi.md`

## 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).
- [x] 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.

Models:

- text models: @ArthurZucker @younesbelkada @fxmarty 

Library:

- trainer: @muellerzr and @SunMarc

Integrations:

- Big Model Inference: @SunMarc

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)

Cc PyTorch Distributed Team:
@gnadathur @wanchaol @fduwjj @wz337 @wconstab @tianyu-l


Comments:
Cc: @lessw2020 @HamidShojanazeri
Cc: @jerryzh168 

Also we would place 
```python
    device = torch.device(f"cuda:{rank}")
    dist.init_process_group("nccl", device_id=device)
    device_mesh = dist.init_device_mesh("cuda", (world_size,))

    with device:
        model = AutoModelForCausalLM.from_pretrained(
            model_id,
        )
```
potentially inside `from_pretrained`, this way we automatically set things for users, with them only having to change the way they call their script (with torchrun or not) 

WDYT?

We will also add this to the `configuration_llama.py`:
```python
DEFAULT_TP_PLAN = {
    "model.layer.*.q_proj":"column_parallel"
}
```
this way we can init lama config with it! 


cc @SunMarc as per our discussion!

Thanks @ArthurZucker @SunMarc for your review.

Good idea to move `tp_plan` to `configuration_llama.py`. I made that change in the current version. Indeed `modeling_llama.py` becomes cleaner now, and it will not have dependency on `torch.distributed...`. 

Since the top-level model can have different `base_model_prefix` (`"model"`, `"transformer"`, etc), I could not directly put FQNs like `"model.layers.*"` in `configuration_llama.py`. So I prepare a `_base_model_tp_plan` starting with `"layers.*"` facing the base model `LlamaModel`. This works because `model.tensor_parallel()` searches for and applies `_tp_plan` recursively. That way we can serve all top models regardless of what prefix name they give to `LlamaModel`, and we can also serve a bare-metal `LlamaModel`.

I also made the TP styles "type neutral" in the config as you suggested. Translation to torch TP types is done only when applying TP.

Re moving distributed / device_mesh init code inside `from_pretrained`:

It is a good direction in making the workflow even simpler for users. As in, wouldn't it be nice if the model is distributed after it is created? I think a UI design would be helpful as to how users express the intention of "please create the model in a distributed manner".

For example, would `from_pretrained` accepts a related argument, like:
```
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_mesh=device_mesh,
)
```
If that kwarg is given, I think the intention is clear enough, and we can call `tensor_parallel()` for the user. Or even better, initialize partial weights only -- this is supported by the proposed torch TP too, but it will require slightly more work in integration with the weight loading part of `from_pretained`. I can try making a demo in a next PR if you are interested.

Re launcher: I might slightly prefer the above "explicit expression" to "implicit dependency" on `torchrun`, for keeping HF transformers "launcher neutral", as some users may use SLURM, `mp.spawn`, etc. Another reason was that `init_process_group` and `init_device_mesh` can take some kwargs (such as timeout, options, etc), which might be a bit big for `from_pretrained` to absorb if it wants to call them inside.

Regardless, I agree that it would be very nice if we can create distributed models out of `from_pretrained" directly, and I am really excite to help that way :) 

Summarizing my discussion with @kmehant re collaboration plan between this PR and https://github.com/huggingface/transformers/pull/34194:
- Both PRs agree on the interest in having PyTorch TP in transformers :)
- We agree that the integration can hopefully happen at a broad base so that it can scale to a set of models, thus we agree on picking `PretrainedModel` as an integration point.
- [#34194](https://github.com/huggingface/transformers/pull/34194) has changes on the trainer side. And there is a corresponding [PR](https://github.com/huggingface/accelerate/pull/3173) for changes on Accelerate side. So those two PRs focus on workflow-level integration. So hopefully this PR can serve as a base for those two PRs and once landed, unblock the other two PRs.

Sounds good, let's focus on this one first then! 

Dumb question: is there a way to bypass the `check_repository_consistency` test? Thanks! @ArthurZucker @SunMarc 

Okay, I just temporarily disabled the "Copied from" comments in some derivative models. Will add them back once we extend TP plans to those models.

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

Thanks @ArthurZucker . I tried that method, but it doesn't seem to work. 
Specifically, the test is checking for consistency between `modeling_gemma.py` and `modular_gemma.py`. I edited neither, but the CI complains that they are not consistent. Here is the [link](https://app.circleci.com/pipelines/github/huggingface/transformers/109404/workflows/7fdcf506-4981-470c-8b3c-17e809855f1a/jobs/1453707).

Ah that is because you probably modified `modeling` directly, when these two have a `modular` file, which has the source of truth!

@ArthurZucker Fixed code consistency issue via `make fix-copies`. Thanks!
All tests seems passed now (except three awaiting approval). I wonder if you have more comments?
Also cc: @kmehant @SunMarc @muellerzr 

I will handle the fix copies afterwards don't worry about it!

Thanks for your comments @ArthurZucker .
The reason I haven't put `init_process_group` or `init_device_mesh` into `from_pretrained` are two folds:

1. I was not sure if I should make `from_pretrained` a collective API.
Both `init_process_group` and `init_device_mesh` have collectivity requirements. If `from_pretrained` calls them inside, it will have that requirement too. So I was hoping to let `init_device_mesh` undertake that requirement and cut `from_pretrained` loose from it. It also means that when user's program hit a hang, the debug / support load is on PyTorch, not on HF. The user could have also called `init_process_group` before calling `from_pretrained`; if we call it a second time, it may error out.

2. I was not sure if the user would like to shard the model over the entire "world".
In a multi-dimensional scenario, the user may want to TP'lize the model only onto a dimension of the world instead of the entire one (which `init_process_group` creates). This implies that the "slice" must come from the user. Of course, that's not the requirement of the example code, but just want to keep ourselves "future proof." 

That said, I agree that we can probably make the program simpler by moving the `.tensor_parallel()` API into `from_pretrained`. Mixing with 1 and 2, how about the following design:
```
tp_mesh = DeviceMesh(...)  # tp_mesh can be smaller than world
model = AutoModelForCausalLM.from_pretrained(
    model_id, tp_mesh=tp_mesh
)
```
This will achieve both 1 and 2:
- `DeviceMesh(...)` will undertake the collectivity requirement, `from_pretrained` will be free from it.
- User can specify which dimension of the world to apply TP via `tp_mesh`.

Chatted with @SunMarc offline. Will move the `DeviceMesh` part in.

Hi @ArthurZucker @SunMarc 
I made the three changes suggested by @ArthurZucker in the comment above.
1. `init_device_mesh` is now moved into `from_pretrained`. When user sets argument `tp_plan="auto"`, the creation will be triggered and `from_pretrained` will return a TP'lized model. That is, the following UI is now supported:
```
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    tp_plan="auto",
)
```
2. `_tp_plan` is now a cls attr of `PreTrainedModel` and I moved `self._tp_plan = config.base_model_tp_plan` from each model's code surface to `PreTrainedModel.post_init()`, so that it is hidden / cleaner.
3. Removed `base_model_tp_plan` from `to_dict` and `to_diff_dict` so that it won't be serialized / saved atm.

Would appreciate your review. Thanks!

And this is benchmarking per @SunMarc suggestion:
![Meta-Llama-3-8B-Instruct, seqlen = 512, python, w_ compile](https://github.com/user-attachments/assets/c8672aca-8b18-441e-abe3-a50fdb33db11)


@ArthurZucker @SunMarc Wondering if you could take another look to see if I addressed your comments? Thanks very much!

> Hi @ArthurZucker @SunMarc I made the three changes suggested by @ArthurZucker in the comment above.
> 
> 1. `init_device_mesh` is now moved into `from_pretrained`. When user sets argument `tp_plan="auto"`, the creation will be triggered and `from_pretrained` will return a TP'lized model. That is, the following UI is now supported:
> 
> ```
> model = AutoModelForCausalLM.from_pretrained(
>     model_id,
>     tp_plan="auto",
> )
> ```
> 
> 2. `_tp_plan` is now a cls attr of `PreTrainedModel` and I moved `self._tp_plan = config.base_model_tp_plan` from each model's code surface to `PreTrainedModel.post_init()`, so that it is hidden / cleaner.
> 3. Removed `base_model_tp_plan` from `to_dict` and `to_diff_dict` so that it won't be serialized / saved atm.
> 
> Would appreciate your review. Thanks!

Sounds absolutely great! :hugs:

@kwen2501 @ArthurZucker 
Do we want to address this minor support - https://github.com/huggingface/transformers/pull/34184#discussion_r1827185086 in this PR? This will unblock accelerate PR. 

@kmehant Thanks for the suggestion. Understood. Agree with the need for a programatic approach.
For the `.tensor_parallel` method itself, it may not be very easy to return a True/False because `nn.Module.apply` does not return anything than the module itself.
Workflow-wise, I wonder if it would be clearer / safer to do:
```
if model.has_tp_plan():
    model.tensor_parallel()
```
If that addresses your need, I am happy to add a `has_tp_plan` util. 

> I love it! TBH We can almost merge as is, the main issue / next PR is adding tests. We need to make sure there are no weird interactions! 🤗 Let's add the graph to the doc somewhere as well !

Good idea. Just looked [here](https://huggingface.co/docs/transformers/testing#distributed-training), seems easy to follow the example. 

https://github.com/huggingface/transformers/pull/34184#issuecomment-2475352157

@kwen2501 Absolutely, this will be helpful, Thanks

Sounds good!

@ArthurZucker Thanks again for the suggestions.
1. A test has been added at `tests/tp/test_tp.py`, run command:
```
CUDA_VISIBLE_DEVICES=0,1 pytest -sv tests/tp/test_tp.py
```
2. A new doc has been added: `perf_infer_gpu_multi.md`, titled "Multi-GPU inference".
A use example and the benchmark figure have been added to the doc.
@ArthurZucker can you please help merge this [image upload](https://huggingface.co/datasets/huggingface/documentation-images/discussions/392) for doc? Thanks!
3. @kmehant The `.has_tp_plan` property has been added.

Merged! let me review one last time 🤗 

@ArthurZucker CI is all green now. (Previous failure seems to come from infra instability). Is there anything else I should add before merge? Thanks!

Nope merging! 🤗 

I wanted to check if `torch._C.get_device_accelerator` had been available for a while or not! 


SUper kudos for the PR! 🚀 

Hi @ArthurZucker and @kwen2501 - it seems this PR now requires that users have torch 2.5+, is that correct?  When trying to run the DeepSpeed unit tests with the latest transformers and torch 2.3 I'm getting errors around not being able to import from `torch.distributed.tensor`:

```
from torch.distributed.tensor import Replicate
```

The specific error:

```
E   RuntimeError: Failed to import transformers.models.vision_encoder_decoder.modeling_vision_encoder_decoder because of the following error (look up to see its traceback):
E   Failed to import transformers.generation.utils because of the following error (look up to see its traceback):
E   cannot import name 'Replicate' from 'torch.distributed.tensor' (/usr/local/lib/python3.10/dist-packages/torch/distributed/tensor/__init__.py)
```

Opened #34795 to track if that's better visibility wise.

@ArthurZucker @kwen2501 Planning for a HF blog? I am in!

@loadams Thanks for reporting it and sorry about the break. 
I think that's because `Replicate` was previously in `torch.distributed._tensor`.
Let me open a PR to build backward compat support.
Cc: @wz337

Sure! Won't have bandwidth to start a blog but feel free to do it! 🤗