Repository: pytorch
Issue: Use Dtensor and ColwiseParallel api to do Tensor Parallel for llama2 model but got communication error.
Labels: oncall: distributed,triaged,module: dtensor

### 🐛 Describe the bug

I read the test_transformer_training example in [pytorch/test/distributed/tensor/parallel/test_tp_examples.py,](https://github.com/pytorch/pytorch/blob/main/test/distributed/tensor/parallel/test_tp_examples.py#L164C9-L164C34) and I think it's really awesome. 
Then I use it to do tensor parallel for Llama2 model and now I may met a communication error.
Here is my llama2 model code, simpily import them from transformers, which version is 4.37.2:

```
from transformers.models.llama.modeling_llama import LlamaForCausalLM
from transformers.models.llama.configuration_llama import LlamaConfig
config = LlamaConfig(
        hidden_size=144,
        intermediate_size=864,
        num_hidden_layers=12,
        num_attention_heads=12,
        max_position_embeddings=512,
    )
    model = LlamaForCausalLM(config)
```
    

And here is my tensor parallel code:


```

from torch.distributed.tensor.parallel import ColwiseParallel, RowwiseParallel, PrepareModuleInput
from torch.distributed._tensor import DeviceMesh, DTensor, Replicate, Shard, Placement
from copy import deepcopy
import torch
import torch.distributed as dist

NUM_DEVICES = 4
if torch.cuda.is_available() and torch.cuda.device_count() > 1:
    NUM_DEVICES = min(NUM_DEVICES, torch.cuda.device_count())

def get_tensor_sharded_model(model, use_ddp):

    if use_ddp:
        model = model.module.model 
    device = "cuda" if torch.cuda.is_available() else "cpu" 
    device_mesh = DeviceMesh(device, torch.arange(0, NUM_DEVICES))
    coordiator = device_mesh.get_coordinate()

    # Parallelize the embedding submodules.
    
    parallelize_module(model.model.embed_tokens, device_mesh, ColwiseParallel(output_layouts=Replicate()))

    # Parallelize the attention and feed forward submodules.
    for layer in model.model.layers:
        
        layer_parallelize_plan = {} 
        layer_parallelize_plan["self_attn.q_proj"] = ColwiseParallel()
        layer_parallelize_plan["self_attn.k_proj"] = ColwiseParallel()
        layer_parallelize_plan["self_attn.v_proj"] = ColwiseParallel()
        layer_parallelize_plan["self_attn.o_proj"] = RowwiseParallel()

        layer_parallelize_plan["mlp.gate_proj"] = ColwiseParallel()
        layer_parallelize_plan["mlp.up_proj"] = ColwiseParallel()
        layer_parallelize_plan["mlp.down_proj"] = RowwiseParallel()

        parallelize_module(layer, device_mesh, layer_parallelize_plan)

    # Parallelize the output submodule. 
    # output = nn.Linear(args.dim, args.vocab_size, bias=False)
    output_parallelize_plan = RowwiseParallel(input_layouts=Replicate())
    parallelize_module(model.lm_head, device_mesh, output_parallelize_plan)

    # Manually adjust the number of heads after sharding the self attention modules.
    for layer in model.model.layers:
        assert model.model.config.num_attention_heads % dist.get_world_size() == 0
        # For llama2 models, your should adjust the number of heads separately.
        layer.self_attn.num_heads = model.model.config.num_attention_heads // dist.get_world_size()
        layer.self_attn.num_key_value_heads = model.model.config.num_key_value_heads // dist.get_world_size()
        layer.self_attn.hidden_size = model.model.config.hidden_size // dist.get_world_size()

    # Manually set output.weight so that parameters and gradients are shared.
    model.lm_head = model.model.embed_tokens.weight
```

And when I do the forward of the llama2 model after tensor shard, it is stucked in the Embedding function which is before attention in llama2 model. Seems like one rank is stucked that I can't go on, and the other three ranks are keep running. I tried it many times and it is always stucked in the same place. It doesn't have any error information or any other tints.
![img_v3_027q_1540648b-71e8-432f-ba94-c41ccaaed06g](https://github.com/pytorch/pytorch/assets/32676639/035de077-778c-4476-a57e-94df491de7ba)

I have excuted your examples in test_tp_examples.py with success.
![img_v3_027o_523b7183-d903-4e4d-8c53-c4ca6661ceag](https://github.com/pytorch/pytorch/assets/32676639/ef965146-88ea-4137-84f5-6475181921f5)

Could anybody help me? I really like the ColwiseParallel and the RowwiseParallel design. They are really fantastic and so incredible that I can write my tensor parallel code so concise, I really want to use them successfully.



### Versions

transformer: 4.37.2

pip list
Package                 Version            Editable project location
----------------------- ------------------ -------------------------
absl-py                 2.1.0
accelerate              0.26.1
aiohttp                 3.9.3
aiosignal               1.3.1
astunparse              1.6.3
attrs                   23.2.0
blessed                 1.20.0
cachetools              5.3.2
certifi                 2024.2.2
charset-normalizer      3.3.2
datasets                2.16.1
dill                    0.3.7
expecttest              0.2.1
filelock                3.13.1
frozenlist              1.4.1
fsspec                  2023.10.0
google-auth             2.26.2
google-auth-oauthlib    1.2.0
gpustat                 1.1.1
grpcio                  1.60.0
huggingface-hub         0.20.3
hypothesis              6.97.4
idna                    3.6
Jinja2                  3.1.3
lark                    1.1.9
Markdown                3.5.2
MarkupSafe              2.1.5
mpmath                  1.3.0
multidict               6.0.5
multiprocess            0.70.15
networkx                3.2.1
numpy                   1.26.3
nvidia-ml-py            12.535.133
oauthlib                3.2.2
optree                  0.10.0
packaging               23.2
pandas                  2.2.0
pip                     23.3.1
prettytable             3.9.0
protobuf                4.23.4
psutil                  5.9.8
pyarrow                 15.0.0
pyarrow-hotfix          0.6
pyasn1                  0.5.1
pyasn1-modules          0.3.0
python-dateutil         2.8.2
pytorch-triton          3.0.0+901819d2b6
pytz                    2024.1
PyYAML                  6.0.1
regex                   2023.12.25
requests                2.31.0
requests-oauthlib       1.3.1
rsa                     4.9
safetensors             0.4.2
sentencepiece           0.1.99
setuptools              68.2.2
six                     1.16.0
sortedcontainers        2.4.0
sympy                   1.12
tensorboard             2.15.1
tensorboard-data-server 0.7.2
tokenizers              0.15.1
torch                   2.3.0a0+gitfcf22a8 /home/wangbinluo/pytorch
tqdm                    4.66.1
transformers            4.37.2
types-dataclasses       0.6.6
typing_extensions       4.9.0
tzdata                  2023.4
urllib3                 2.2.0
wcwidth                 0.2.13
Werkzeug                3.0.1
wheel                   0.41.2
xxhash                  3.4.1
yarl                    1.9.4

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @XilunWu @wanchaol @fduwjj @wz337 @tianyu-l @wconstab @yf225