Repository: pytorch
Issue: The forward accuracy cannot be aligned when using  DTensor and ColwiseParallel, RowwiseParallel for TP training on the Llama2
Labels: oncall: distributed,module: dtensor

### 🐛 Describe the bug

Hi, I found that the forward accuracy cannot be aligned completely when using your 'DTensor' and 'ColwiseParallel', 'RowwiseParallel', and 'Replicated' api for TP training on the Llama2. 

I referred your example: https://github.com/pytorch/pytorch/blob/main/test/distributed/tensor/parallel/test_tp_examples.py#L175C9-L262C1. I parallelize the embedding submodules of Llama2 by 'Replicate()', parallelize the attention and feed forward by 'ColwiseParallel' and 'RowwiseParallel', and Replicate the model.lm_head for the output, and manually adjust the number of heads and the hidden size.

The following is my code:

```
def get_tensor_sharded_model(model, use_ddp):

    device = "cuda" if torch.cuda.is_available() else "cpu" 
    device_mesh = DeviceMesh(device, torch.arange(0, NUM_DEVICES))
    # 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_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()
```

And here is my loss comparison  between model and model_tp:
![img_v3_0284_103fb69e-539e-427f-8d3f-29e7981c5f7g](https://github.com/pytorch/pytorch/assets/32676639/7fb6593a-6f61-4d88-8ffe-78bf742716c6)
and my comparison code, just like you used in your example:
![img_v3_0284_72d46648-041d-474f-9004-909adbe6764g](https://github.com/pytorch/pytorch/assets/32676639/18510bc6-8466-4004-a9aa-556cb461d0a4)

It said I have 29% elements mismatched in my first step loss comparison!

The gradient of model parameters in backward cannot be aligned too.

![img_v3_0284_c152df36-41e2-4b80-8dc2-70d5a3ae96bg](https://github.com/pytorch/pytorch/assets/32676639/b945eb38-a4ab-47d9-a5c2-04986e4da807)


I feel very confused and have invested the whole process, like whether the input is consistent across each rank(it is), and whether the dropout in the llama2 model could introduce randomness (Actually llama2 does not have dropout). Now I really don't know why. Could anybody help me to check this? Should I add some all_reduce for the loss or something?


### Versions

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