[Bug-Qwen2VL] Error when calling generate with fsdp2
Author: mantle2048
Pull request: no
Labels: name: bug

name: VLM

I have split Qwen2VL-2B using fsdp2 across 8 GPUs. Following the [official example](https://github.com/huggingface/transformers/blob/main/tests/generation/test_fsdp.py#L81-L101), I call the generate function, but encounter the following error:

```markdown
[rank7]:   File "/software/mamba/envs/mm/lib/python3.11/site-packages/transformers/models/qwen2_vl/modeling_qwen2_vl.py", line 1791, in prepare_inputs_for_generation
[rank7]:     input_ids = input_ids[:, cache_position]
[rank7]:                 ~~~~~~~~~^^^^^^^^^^^^^^^^^^^
[rank7]: RuntimeError: CUDA error: device-side assert triggered
[rank7]: Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions
```

Upon inspecting the source code of `modeling_qwen2_vl.py`, I found that the `cache_position` in the[ line 1735](https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py#L1735) `input_ids = input_ids[:, cache_position]` is out of bounds.

It seems that the generation has already stopped and `input_ids` is no longer being extended, but `cache_position` is still increasing. I am unsure of the deeper cause of this issue.

A quick fix I found is to add `cache_position[0] = min(cache_position[0], input_ids.shape[1] - 1)` above this line.

I’d like to ask for your help—does this seem like a viable fix? If not, is there a way to make generate work correctly under fsdp2 for Qwen2VL without modifying the source code?

Here is a minimal reproduction script (transformers.__version__
==4.47.0). After the above fix, this script should run correctly on 8 GPUs.

```python
import torch
from transformers import Qwen2VLForConditionalGeneration, AutoProcessor

from qwen_vl_utils import process_vision_info
import torch
import torch.distributed
from torch.distributed._composable.fsdp import fully_shard, register_fsdp_forward_method
from torch.distributed.device_mesh import init_device_mesh

from transformers.models.qwen2_vl.modeling_qwen2_vl import Qwen2VLDecoderLayer

queries = [
    "What is the main object shown in the picture?",
    "What type of plants are present in the image?",
    "What items can be seen on the table?",
    "Is the view outside the window sunny or cloudy?",
    "Can you spot any drinks in the image?",
    "Are there any decorations or paintings on the wall?",
    "What is the color and material of the chair?",
    "Is the scene set indoors or outdoors?"
]
data = []
for query in queries:
    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "image",
                    "image": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg",
                },
                {"type": "text", "text": query},
            ],
        }
    ]
    data.append(messages)


def fsdp2_generate():
    torch.cuda.set_device(device := torch.device(rank := torch.distributed.get_rank()))

    model = Qwen2VLForConditionalGeneration.from_pretrained(
        "Qwen2/Qwen2-VL-2B-Instruct",
        device_map="cpu"
    )

    model.to(device)

    mesh = init_device_mesh("cuda", (torch.distributed.get_world_size(),))
    for submodule in model.modules():
        if isinstance(submodule, Qwen2VLDecoderLayer):
            fully_shard(submodule, mesh=mesh)
    fully_shard(model, mesh=mesh)

    register_fsdp_forward_method(model, "generate")

    processor = AutoProcessor.from_pretrained("Qwen2/Qwen2-VL-2B-Instruct")
    messages = data[rank]
    text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    image_inputs, _ = process_vision_info(messages)
    batch = processor(
        text=[text],
        images=image_inputs,
        return_tensors="pt",
    ).to(device)

    query_responses = model.generate(
        **batch,
        max_new_tokens=128,
    )
    context_length = batch.input_ids.shape[1]
    responses = query_responses[:, context_length:]
    response_texts = processor.batch_decode(
        responses, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )

    rank = torch.distributed.get_rank()
    print()
    print(f"Rank {rank}: {response_texts}")
    print()


if __name__ == "__main__":
    torch.distributed.init_process_group(backend='nccl', world_size=torch.cuda.device_count())
    fsdp2_generate()
    torch.distributed.destroy_process_group()
```

### Who can help?

@amyeroberts @qubvel @zucchini-nlp 

### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

- [x] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

```python
import torch
from transformers import Qwen2VLForConditionalGeneration, AutoProcessor

from qwen_vl_utils import process_vision_info
import torch
import torch.distributed
from torch.distributed._composable.fsdp import fully_shard, register_fsdp_forward_method
from torch.distributed.device_mesh import init_device_mesh

from transformers.models.qwen2_vl.modeling_qwen2_vl import Qwen2VLDecoderLayer

queries = [
    "What is the main object shown in the picture?",
    "What type of plants are present in the image?",
    "What items can be seen on the table?",
    "Is the view outside the window sunny or cloudy?",
    "Can you spot any drinks in the image?",
    "Are there any decorations or paintings on the wall?",
    "What is the color and material of the chair?",
    "Is the scene set indoors or outdoors?"
]
data = []
for query in queries:
    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "image",
                    "image": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg",
                },
                {"type": "text", "text": query},
            ],
        }
    ]
    data.append(messages)


def fsdp2_generate():
    torch.cuda.set_device(device := torch.device(rank := torch.distributed.get_rank()))

    model = Qwen2VLForConditionalGeneration.from_pretrained(
        "Qwen2/Qwen2-VL-2B-Instruct",
        device_map="cpu"
    )

    model.to(device)

    mesh = init_device_mesh("cuda", (torch.distributed.get_world_size(),))
    for submodule in model.modules():
        if isinstance(submodule, Qwen2VLDecoderLayer):
            fully_shard(submodule, mesh=mesh)
    fully_shard(model, mesh=mesh)

    register_fsdp_forward_method(model, "generate")

    processor = AutoProcessor.from_pretrained("Qwen2/Qwen2-VL-2B-Instruct")
    messages = data[rank]
    text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    image_inputs, _ = process_vision_info(messages)
    batch = processor(
        text=[text],
        images=image_inputs,
        return_tensors="pt",
    ).to(device)

    query_responses = model.generate(
        **batch,
        max_new_tokens=128,
    )
    context_length = batch.input_ids.shape[1]
    responses = query_responses[:, context_length:]
    response_texts = processor.batch_decode(
        responses, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )

    rank = torch.distributed.get_rank()
    print()
    print(f"Rank {rank}: {response_texts}")
    print()


if __name__ == "__main__":
    torch.distributed.init_process_group(backend='nccl', world_size=torch.cuda.device_count())
    fsdp2_generate()
    torch.distributed.destroy_process_group()
```

### Expected behavior

All ranks can generate correct outputs.

Comments:
@mantle2048 hey! That seems to be fixed by https://github.com/huggingface/transformers/pull/35893, which is still not merged into `main`. can you try from the PR branch if that solves your issue?

@zucchini-nlp Hi! I tried this PR and now the multi gpu generation is working very well. Look forward to its merger soon😊