Multi GPU generate with llama shape error
Author: dakinggg
Pull request: no
Labels: name: WIP

name: bug

name: Generation

### System Info

- `transformers` version: 4.44.0
- Platform: Linux-5.15.0-91-generic-x86_64-with-glibc2.31
- Python version: 3.10.12
- Huggingface_hub version: 0.23.4
- Safetensors version: 0.4.2
- Accelerate version: 0.25.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.2.1+cu121 (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?: yes, running with `composer` as the distributed launcher
- Using GPU in script?: yes
- GPU type: NVIDIA A100-SXM4-40GB

### Who can help?

@gante @ArthurZucker 

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Run llama with `synced_gpus=True` and an attention mask. This worked fine on transformers 4.40.2 (and 4.41.x), but no longer works. The use of Composer for the dist stuff is just convenience, shouldn't affect anything to swap in a different distributed launcher, etc.

```python3
import transformers
import torch
from composer.utils import dist

def main():
    dist.initialize_dist('gpu')

    name = 'meta-llama/Meta-Llama-3-8B-Instruct'
    tokenizer = transformers.AutoTokenizer.from_pretrained(name)
    pad_token_id = tokenizer.eos_token_id
    model = transformers.AutoModelForCausalLM.from_pretrained(name)

    rank = dist.get_global_rank()

    model.to(f'cuda:{rank}')

    if dist.get_global_rank() == 0:
        content = 'Write one short sentence.'
    else:
        content = 'Write one long paragraph.'

    messages = [
        {
            'role': 'user',
            'content': content,
        }
    ]

    tokenized_messages = tokenizer.apply_chat_template(messages, return_tensors='pt')

    padded_messages = torch.cat(
        [
            torch.LongTensor((4096 - 20) * [pad_token_id]),
            tokenized_messages[0],  # [seq]
        ],
        dim=0,
    )
    padded_messages = padded_messages.unsqueeze(0)
    padded_messages = padded_messages.to(f'cuda:{rank}')
    attention_mask = ~(padded_messages == pad_token_id)
    attention_mask = attention_mask.to(f'cuda:{rank}')
    output = model.generate(input_ids=padded_messages, attention_mask=attention_mask, synced_gpus=True, max_new_tokens=200)

    print(tokenizer.decode(output[0]))

if __name__ == '__main__':
    main()
```

This results in 
```
Traceback (most recent call last):
  File "/mnt/workdisk/danielking/github/multi-gpu.py", line 47, in <module>
    main()
  File "/mnt/workdisk/danielking/github/multi-gpu.py", line 42, in main
    output = model.generate(input_ids=padded_messages, attention_mask=attention_mask, synced_gpus=True, max_new_tokens=200)
  File "/mnt/workdisk/danielking/miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/mnt/workdisk/danielking/miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/transformers/generation/utils.py", line 2024, in generate
    result = self._sample(
  File "/mnt/workdisk/danielking/miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/transformers/generation/utils.py", line 2982, in _sample
    outputs = self(**model_inputs, return_dict=True)
  File "/mnt/workdisk/danielking/miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/mnt/workdisk/danielking/miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/mnt/workdisk/danielking/miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 1189, in forward
    outputs = self.model(
  File "/mnt/workdisk/danielking/miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/mnt/workdisk/danielking/miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/mnt/workdisk/danielking/miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 1001, in forward
    layer_outputs = decoder_layer(
  File "/mnt/workdisk/danielking/miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/mnt/workdisk/danielking/miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/mnt/workdisk/danielking/miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 734, in forward
    hidden_states, self_attn_weights, present_key_value = self.self_attn(
  File "/mnt/workdisk/danielking/miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/mnt/workdisk/danielking/miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/mnt/workdisk/danielking/miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 660, in forward
    attn_output = torch.nn.functional.scaled_dot_product_attention(
RuntimeError: The expanded size of the tensor (4105) must match the existing size (4104) at non-singleton dimension 3.  Target sizes: [1, 32, 1, 4105].  Tensor sizes: [1, 1, 1, 4104]
```

### Expected behavior

Multi GPU generate does not error..

Comments:
Hi @dakinggg 👋 

I tried to reproduce your issue on the latest `main` (commit = `22e6f145259c4e9e156a3f846fad14b1cddfae4f`), and it ran on my end with 2xA100 80GB

Can you confirm that you can reproduce your issue on `main`? If so, we'll need to dig deeper 🤔 

hi @gante, I get the same issue on transformers main

commit for me was `0a7af19f4dc868bafc82f35eb7e8d13bac87a594`

@dakinggg I've also ran against `0a7af19f4dc868bafc82f35eb7e8d13bac87a594`, with no issues on my end 🤔 

Can you confirm that:
1. you don't have another version of `transformers` installed locally? (`transformers.__version__` should return `'4.45.0.dev0'`)
2. you have accelerate v0.33
3. you have torch 2.1.2 (`composer` forces `torch==2.1.2` on my end)




1. confirmed
2. confirmed
3. 2.3.1, the last release of composer supports 2.1, 2.2, 2.3, and we just released a new version this morning that shifts to 2.2, 2.3, 2.4.

let me make a new environment though and retry

Just tried again and same error 🤔 

for completeness
```
Package                  Version        Editable project location
------------------------ -------------- --------------------------------------------
accelerate               0.33.0
anyio                    4.4.0
argcomplete              3.5.0
arrow                    1.3.0
asttokens                2.4.1
backoff                  2.2.1
certifi                  2024.7.4
charset-normalizer       3.3.2
coolname                 2.2.0
decorator                5.1.1
exceptiongroup           1.2.2
executing                2.0.1
filelock                 3.15.4
fsspec                   2024.6.1
gql                      3.5.0
graphql-core             3.2.3
huggingface-hub          0.24.6
idna                     3.8
importlib_metadata       8.4.0
ipython                  8.26.0
jedi                     0.19.1
Jinja2                   3.1.4
lightning-utilities      0.11.6
markdown-it-py           3.0.0
MarkupSafe               2.1.5
matplotlib-inline        0.1.7
mdurl                    0.1.2
mosaicml                 0.24.0
mosaicml-cli             0.6.41
mpmath                   1.3.0
multidict                6.0.5
networkx                 3.3
numpy                    1.26.4
nvidia-cublas-cu12       12.1.3.1
nvidia-cuda-cupti-cu12   12.1.105
nvidia-cuda-nvrtc-cu12   12.1.105
nvidia-cuda-runtime-cu12 12.1.105
nvidia-cudnn-cu12        9.1.0.70
nvidia-cufft-cu12        11.0.2.54
nvidia-curand-cu12       10.3.2.106
nvidia-cusolver-cu12     11.4.5.107
nvidia-cusparse-cu12     12.1.0.106
nvidia-nccl-cu12         2.20.5
nvidia-nvjitlink-cu12    12.6.20
nvidia-nvtx-cu12         12.1.105
packaging                24.1
parso                    0.8.4
pexpect                  4.9.0
pillow                   10.4.0
pip                      24.2
prompt-toolkit           3.0.36
protobuf                 5.27.3
psutil                   6.0.0
ptyprocess               0.7.0
pure_eval                0.2.3
py-cpuinfo               9.0.0
Pygments                 2.18.0
python-dateutil          2.9.0.post0
pytorch-ranger           0.1.1
PyYAML                   6.0.2
questionary              2.0.1
regex                    2024.7.24
requests                 2.32.3
rich                     13.8.0
ruamel.yaml              0.18.6
ruamel.yaml.clib         0.2.8
safetensors              0.4.4
setuptools               72.1.0
six                      1.16.0
sniffio                  1.3.1
stack-data               0.6.3
sympy                    1.13.2
tabulate                 0.9.0
termcolor                2.4.0
tokenizers               0.19.1
torch                    2.4.0
torch-optimizer          0.3.0
torchmetrics             1.4.0.post0
torchvision              0.19.0
tqdm                     4.66.5
traitlets                5.14.3
transformers             4.45.0.dev0    /mnt/workdisk/danielking/github/transformers
triton                   3.0.0
types-python-dateutil    2.9.0.20240821
typing_extensions        4.12.2
urllib3                  2.2.2
validators               0.33.0
wcwidth                  0.2.13
websockets               11.0.3
wheel                    0.43.0
yarl                     1.9.4
zipp                     3.20.0
```

And also to confirm, if, in that same environment, i downgrade transformers (`pip install 'transformers<4.42'`), it works as expected

also removed the composer dependency just to be sure and it still errors in the same way

For completeness, I downgraded my torch to the version you used (2.1.2), and get the same error

I bisected it down to this PR: https://github.com/huggingface/transformers/pull/30536/files

@dakinggg what we're seeing (failing on your end, not failing on mine) probably can be explained as a version-dependent bug. I've created a fresh venv and deleted `pip` cache, resulting in the same versions and same outcomes. I am noticing, however, that I get different versions from yours (e.g. `composer==0.18.2`, `torch==2.1.2`). This seems to be due to my machine having python 3.8.

I don't have permissions on my multi-gpu machine to try different python versions -- do you have on yours? If so, can you try python 3.8? If not, let me know, I'll try to request a machine for my to debug this on :)

ah, hold up, I found a shared multi-gpu machine with python 3.10! Testing things up

@dakinggg your test script passing on our python 3.10 machine, with `composer==0.24.0`, torch `2.4.0`, ... 🤔. This was a machine where I had a fresh environment on CUDA 12.0  

To attempt rule out environment issues: if you have access to a different machine, ideally without ML dependencies installed, can you attempt to reproduce it there?

(I'm afraid I'm not sure how to help without being able to reproduce the issue)


Super strange...was hoping that getting the right PR would be enough of a pointer in the right direction. Are there any parts of the generate loop that depend on package versions that you know of?

btw, you shouldn't need any extra permissions to try a new python version. I usually use mamba package manager to get environments with specific python versions. I've been creating a fresh new environment for these tests, tried both python 3.10 and 3.11. Thinking what else I can try...

Update: after offline conversation with @dakinggg, we've found how to reproduce -- the command has to be launched with `composer` (e.g. `composer foo.py`)

*composer not compose

and also here is a version of the script that doesn't rely on a distributed launcher and can be run directly with `python`

```
import transformers
import os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp

def run(rank, size):
    # dist.initialize_dist('gpu')

    name = 'meta-llama/Meta-Llama-3-8B-Instruct'
    tokenizer = transformers.AutoTokenizer.from_pretrained(name)
    pad_token_id = tokenizer.eos_token_id
    model = transformers.AutoModelForCausalLM.from_pretrained(name)

    # rank = dist.get_global_rank()

    model.to(f'cuda:{rank}')

    if rank == 0:
        content = 'Write one short sentence.'
    else:
        content = 'Write one long paragraph.'

    messages = [
        {
            'role': 'user',
            'content': content,
        }
    ]

    tokenized_messages = tokenizer.apply_chat_template(messages, return_tensors='pt')

    padded_messages = torch.cat(
        [
            torch.LongTensor((4096 - 20) * [pad_token_id]),
            tokenized_messages[0],  # [seq]
        ],
        dim=0,
    )
    padded_messages = padded_messages.unsqueeze(0)
    padded_messages = padded_messages.to(f'cuda:{rank}')
    attention_mask = ~(padded_messages == pad_token_id)
    attention_mask = attention_mask.to(f'cuda:{rank}')
    output = model.generate(input_ids=padded_messages, attention_mask=attention_mask, synced_gpus=True, max_new_tokens=200)

    print(tokenizer.decode(output[0]))

def init_process(rank, size, fn, backend='gloo'):
    """ Initialize the distributed environment. """
    os.environ['MASTER_ADDR'] = '127.0.0.1'
    os.environ['MASTER_PORT'] = '29500'
    dist.init_process_group(backend, rank=rank, world_size=size)
    fn(rank, size)

if __name__ == "__main__":
    size = 2
    processes = []
    mp.set_start_method("spawn")
    for rank in range(size):
        p = mp.Process(target=init_process, args=(rank, size, run))
        p.start()
        processes.append(p)

    for p in processes:
        p.join()
```

Okay, nailed the root causes with `synced_gpus`:
1. `cache_position` stops being updated when generation finishes in a given device, causing the cache indexing errors
2. if we continue updating `cache_position`, then slicing `input_ids` gets out of bounds for the dummy computations

#33252 would fix it, but it conflicts with a larger refactor around `prepare_inputs_for_generation`. As such, I've made fixing this issue part of the refactor goals 🤗 See #32685 

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.

Not stale I think?

#34095 fixes this issue on the vast majority of models 🤗