[model] fix: refactor qwen2vl patches & support no-image input for fsdp (verl-project#3496)

### What does this PR do?

This PR tries to fix verl-project#3491 

### Checklist Before Starting

- [x] Search for similar PRs. Paste at least one query link here: ...
- [x] Format the PR title as `[{modules}] {type}: {description}` (This
will be checked by the CI)
- `{modules}` include `fsdp`, `megatron`, `sglang`, `vllm`, `rollout`,
`trainer`, `ci`, `training_utils`, `recipe`, `hardware`, `deployment`,
`ray`, `worker`, `single_controller`, `misc`, `perf`, `model`, `algo`,
`env`, `tool`, `ckpt`, `doc`, `data`
- If this PR involves multiple modules, separate them with `,` like
`[megatron, fsdp, doc]`
  - `{type}` is in `feat`, `fix`, `refactor`, `chore`, `test`
- If this PR breaks any API (CLI arguments, config, function signature,
etc.), add `[BREAKING]` to the beginning of the title.
  - Example: `[BREAKING][fsdp, megatron] feat: dynamic batching`

### Test

Tested with [latest
transformers](https://github.com/huggingface/transformers/tree/6e50a8afb2540ac1acaa4b62cf1dd5f1170f6511)

<img width="2448" height="540" alt="image"
src="https://github.com/user-attachments/assets/06d40f40-572c-4454-8e08-115857f61f21"
/>
<img width="2796" height="1394" alt="image"
src="https://github.com/user-attachments/assets/17489b9c-e376-46e3-80d8-71106d304077"
/>
<img width="2098" height="744" alt="image"
src="https://github.com/user-attachments/assets/8c7f736d-bf09-4ba9-9cf4-0d56e367c526"
/>

### API and Usage Example

> Demonstrate how the API changes if any, and provide usage example(s)
if possible.

```python
# Add code snippet or script demonstrating how to use this
```

### Design & Code Changes

#### ⚠️ Breaking

We adopt a new format for Qwen2VL's position ids: (4, batch size, seq
len)

Assuming a vision position ids (mrope) has a shape of (3, batch size,
seq len) and a text position ids (normal rope) has a shape of (1, batch
size, seq len), we concatenate both to obtain the final position ids.

This aligns with the implementation in the Transformers >= 4.54.0 🤗

https://github.com/huggingface/transformers/blob/v4.54.0/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py#L1469

#### 🎤 New

We have refactored the Qwen2VL and Qwen2.5VL patches, supporting
no-image input for FSDP by introducing fake ViT inputs. We have also
removed some redundant code for better maintainability.

#### 🚨 Changes

We move the ulysses logic into the attention function. So the position
ids will be scattered before the language model part.

### Checklist Before Submitting

> [!IMPORTANT]
> Please check all the following items before requesting a review,
otherwise the reviewer might deprioritize this PR for review.

- [x] Read the [Contribute
Guide](https://github.com/volcengine/verl/blob/main/CONTRIBUTING.md).
- [x] Apply [pre-commit
checks](https://github.com/volcengine/verl/blob/main/CONTRIBUTING.md#code-linting-and-formatting):
`pre-commit install && pre-commit run --all-files --show-diff-on-failure
--color=always`
- [ ] Add / Update [the
documentation](https://github.com/volcengine/verl/tree/main/docs).
- [x] Add unit or end-to-end test(s) to [the CI
workflow](https://github.com/volcengine/verl/tree/main/.github/workflows)
to cover all the code. If not feasible, explain why: ...
- [x] Once your PR is ready for CI, send a message in [the `ci-request`
channel](https://verl-project.slack.com/archives/C091TCESWB1) in [the
`verl` Slack
workspace](https://join.slack.com/t/verl-project/shared_invite/zt-3855yhg8g-CTkqXu~hKojPCmo7k_yXTQ).
(If not accessible, please try [the Feishu group
(飞书群)](https://applink.larkoffice.com/client/chat/chatter/add_by_link?link_token=772jd4f1-cd91-441e-a820-498c6614126a).)

Author: hiyouga

examples/grpo_trainer/run_qwen2_5_vl-7b.sh

@@ -14,6 +14,7 @@ python3 -m verl.trainer.main_ppo \
     actor_rollout_ref.model.path=Qwen/Qwen2.5-VL-7B-Instruct \
     actor_rollout_ref.actor.optim.lr=1e-6 \
     actor_rollout_ref.model.use_remove_padding=True \
+    actor_rollout_ref.model.use_fused_kernels=True \
     actor_rollout_ref.actor.ppo_mini_batch_size=128 \
     actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=10 \
     actor_rollout_ref.actor.use_kl_loss=True \

verl/models/transformers/monkey_patch.py

@@ -15,17 +15,15 @@
 Apply monkey-patch function to models
 """
 
-import importlib.metadata
 import sys
-from functools import lru_cache
 from typing import Optional
 
 import torch
-from packaging import version
 from transformers.modeling_flash_attention_utils import _flash_attention_forward
 from transformers.modeling_utils import PreTrainedModel
 
 from verl.utils.import_utils import is_trl_available
+from verl.utils.transformers_compat import is_transformers_version_in_range
 from verl.utils.ulysses import (
     gather_heads_scatter_seq,
     gather_seq_scatter_heads,
@@ -51,13 +49,19 @@ def _ulysses_flash_attention_forward(
     query_states: torch.Tensor,
     key_states: torch.Tensor,
     value_states: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    query_length: int,
     *args,
     position_ids: Optional[torch.Tensor] = None,
     **kwargs,
 ):
     """Insert all-to-all before and after flash attention.
     DeepSpeed-Ulysses: https://arxiv.org/pdf/2309.14509
 
+    For transformers>=4.55, the flash attention api has changed,
+    we need to pass the query_length after doing ulysses all2all.
+    See https://github.com/huggingface/transformers/issues/40399
+
     Args:
         query_states (torch.Tensor): (batch_size, seqlen/sp_size, nheads, head_dim)
         key_states (torch.Tensor): (batch_size, seqlen/sp_size, nheads_k, head_dim)
@@ -66,64 +70,7 @@ def _ulysses_flash_attention_forward(
 
     Returns:
         torch.Tensor: (batch_size, seqlen/sp_size, nheads, head_dim)
-    """
-    ulysses_sp_size = get_ulysses_sequence_parallel_world_size()
-
-    ########## AlltoAll for Ulysses ##########
-    if ulysses_sp_size > 1:
-        assert position_ids is not None, "position_ids is required for Ulysses sequence parallelism"
-
-        # NOTE: repeat kv heads to be divided by sequence parallel. Instead of repeating nheads_q//nheads_k,
-        # we choose to repeat sp_size//nheads_k, since flash_attention supports MQA/GQA.
-        # For example:
-        # - nheads_k=4, sp=8, repeats=2
-        # - nheads_k=8, sp=8, repeats=1
-        # - nheads_k=16, sp=8, repeats=1
-        repeats = max(ulysses_sp_size // key_states.size(2), 1)
-        key_states = repeat_kv(key_states, repeats)
-        value_states = repeat_kv(value_states, repeats)
-
-        # (bsz, seq_len/n, n_head, head_dim) -> (bsz, seq_len, n_head/n, head_dim)
-        query_states = gather_seq_scatter_heads(query_states, seq_dim=1, head_dim=2)
-        key_states = gather_seq_scatter_heads(key_states, seq_dim=1, head_dim=2)
-        value_states = gather_seq_scatter_heads(value_states, seq_dim=1, head_dim=2)
-
-        # TODO: all_gather position_ids because `prepare_fa2_from_position_ids` needs it, we can eliminate
-        # this all_gather by passing cu_seq_lens_q, cu_seq_lens_k, max_length_k, max_length_q explicitly.
-        # https://github.com/huggingface/transformers/pull/33932
-
-        # (bsz, seq_len/n) -> (bsz, seq_len)
-        position_ids_list = [torch.empty_like(position_ids) for _ in range(ulysses_sp_size)]
-        torch.distributed.all_gather(position_ids_list, position_ids, group=get_ulysses_sequence_parallel_group())
-        position_ids = torch.concat(position_ids_list, dim=-1)
-
-    # (bsz, seq_len, n_head/n, head_dim)
-    attn_output = _flash_attention_forward(
-        query_states, key_states, value_states, *args, position_ids=position_ids, **kwargs
-    )
-
-    ########## AlltoAll for Ulysses ##########
-    if ulysses_sp_size > 1:
-        # (bsz, seq_len, n_head/n, head_dim) -> (bsz, seq_len/n, n_head, head_dim)
-        attn_output = gather_heads_scatter_seq(attn_output, seq_dim=1, head_dim=2)
 
-    return attn_output
-
-
-def _ulysses_flash_attention_forward_transformers_4_55(
-    query_states: torch.Tensor,
-    key_states: torch.Tensor,
-    value_states: torch.Tensor,
-    attention_mask: Optional[torch.Tensor],
-    query_length: int,
-    *args,
-    position_ids: Optional[torch.Tensor] = None,
-    **kwargs,
-):
-    """For transformers>=4.55, the flash attention api has changed,
-    we need to pass the query_length after doing ulysses alltoall.
-
-    See https://github.com/huggingface/transformers/issues/40399
     """
     ulysses_sp_size = get_ulysses_sequence_parallel_world_size()
 
@@ -178,6 +125,7 @@ def patch_vlm_for_ulysses_input_slicing(model_class: type):
     def _create_ulysses_wrapped_decoder_forward(original_forward):
         def ulysses_wrapped_decoder_forward(self, *args, **kwargs):
             inputs_embeds = kwargs.get("inputs_embeds")
+            position_ids = kwargs.get("position_ids")
             call_kwargs = kwargs.copy()
 
             current_ulysses_sp_size = get_ulysses_sequence_parallel_world_size()
@@ -189,6 +137,7 @@ def ulysses_wrapped_decoder_forward(self, *args, **kwargs):
             )
             if slice_now:
                 call_kwargs["inputs_embeds"] = slice_input_tensor(inputs_embeds, dim=1, padding=False)
+                call_kwargs["position_ids"] = slice_input_tensor(position_ids, dim=-1, padding=False)
                 self._needs_initial_slice = False
             try:
                 return original_forward(self, *args, **call_kwargs)
@@ -225,12 +174,7 @@ def patch_forward_with_backends(
 
     forward_with_torch_backend_function = model.__class__.forward
     forward_with_triton_backend_function = model.__class__.forward
-    if model.config.model_type == "qwen2_5_vl":
-        from verl.models.transformers.qwen2_5_vl import forward_with_torch_backend, forward_with_triton_backend
-
-        forward_with_torch_backend_function = forward_with_torch_backend
-        forward_with_triton_backend_function = forward_with_triton_backend
-    elif model.config.model_type == "qwen2_vl":
+    if model.config.model_type in ["qwen2_5_vl", "qwen2_vl"]:
         from verl.models.transformers.qwen2_vl import forward_with_torch_backend, forward_with_triton_backend
 
         forward_with_torch_backend_function = forward_with_torch_backend
@@ -296,50 +240,70 @@ def state_dict(self, *args, **kwargs):
 
     # TODO: VLM models only, unify monkey patch to LLM models.
     if model.config.model_type == "qwen2_5_vl":
-        if is_transformers_version_in_range(min_version="4.53.0"):
-            from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VLAttention
+        if is_transformers_version_in_range(min_version="4.52.0"):
+            from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
+                Qwen2_5_VLAttention,
+                Qwen2_5_VLForConditionalGeneration,
+                Qwen2_5_VLModel,
+                Qwen2_5_VLTextModel,
+            )
+
+            from verl.models.transformers.qwen2_vl import forward_with_normal_backend, qwen2_vl_base_forward
+
+            Qwen2_5_VLModel.forward = qwen2_vl_base_forward
+            Qwen2_5_VLForConditionalGeneration.forward = forward_with_normal_backend
         else:
             from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
                 Qwen2_5_VLFlashAttention2 as Qwen2_5_VLAttention,
             )
+            from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
+                Qwen2_5_VLForConditionalGeneration,
+            )
+            from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VLModel as Qwen2_5_VLTextModel
+
+            from verl.models.transformers.qwen2_vl import forward_with_normal_backend
+
+            Qwen2_5_VLForConditionalGeneration.forward = forward_with_normal_backend
 
         if use_remove_padding or ulysses_sp_size > 1:
-            from verl.models.transformers.qwen2_vl import ulysses_flash_attn_forward
+            from verl.models.transformers.qwen2_vl import qwen2_vl_attn_forward
 
-            Qwen2_5_VLAttention.forward = ulysses_flash_attn_forward
-            print("Monkey patch FlashAttention2.forward in Qwen2.5VL")
+            Qwen2_5_VLAttention.forward = qwen2_vl_attn_forward
+            print("Monkey patch Qwen2.5VL attention layer")
 
         if ulysses_sp_size > 1:
-            if is_transformers_version_in_range(min_version="4.52.0"):
-                from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VLTextModel
+            patch_vlm_for_ulysses_input_slicing(Qwen2_5_VLTextModel)
 
-                patch_vlm_for_ulysses_input_slicing(Qwen2_5_VLTextModel)
-            else:
-                from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VLModel
+    elif model.config.model_type == "qwen2_vl":
+        if is_transformers_version_in_range(min_version="4.52.0"):
+            from transformers.models.qwen2_vl.modeling_qwen2_vl import (
+                Qwen2VLAttention,
+                Qwen2VLForConditionalGeneration,
+                Qwen2VLModel,
+                Qwen2VLTextModel,
+            )
 
-                patch_vlm_for_ulysses_input_slicing(Qwen2_5_VLModel)
+            from verl.models.transformers.qwen2_vl import forward_with_normal_backend, qwen2_vl_base_forward
 
-    elif model.config.model_type == "qwen2_vl":
-        if is_transformers_version_in_range(min_version="4.53.0"):
-            from transformers.models.qwen2_vl.modeling_qwen2_vl import Qwen2VLAttention
+            Qwen2VLModel.forward = qwen2_vl_base_forward
+            Qwen2VLForConditionalGeneration.forward = forward_with_normal_backend
         else:
             from transformers.models.qwen2_vl.modeling_qwen2_vl import Qwen2VLFlashAttention2 as Qwen2VLAttention
+            from transformers.models.qwen2_vl.modeling_qwen2_vl import Qwen2VLForConditionalGeneration
+            from transformers.models.qwen2_vl.modeling_qwen2_vl import Qwen2VLModel as Qwen2VLTextModel
 
-        if use_remove_padding or ulysses_sp_size > 1:
-            from verl.models.transformers.qwen2_vl import ulysses_flash_attn_forward
+            from verl.models.transformers.qwen2_vl import forward_with_normal_backend
 
-            Qwen2VLAttention.forward = ulysses_flash_attn_forward
-            print("Monkey patch FlashAttention2.forward in Qwen2VL")
+            Qwen2VLForConditionalGeneration.forward = forward_with_normal_backend
 
-        if ulysses_sp_size > 1:
-            if is_transformers_version_in_range(min_version="4.52.0"):
-                from transformers.models.qwen2_vl.modeling_qwen2_vl import Qwen2VLTextModel
+        if use_remove_padding or ulysses_sp_size > 1:
+            from verl.models.transformers.qwen2_vl import qwen2_vl_attn_forward
 
-                patch_vlm_for_ulysses_input_slicing(Qwen2VLTextModel)
-            else:
-                from transformers.models.qwen2_vl.modeling_qwen2_vl import Qwen2VLModel
+            Qwen2VLAttention.forward = qwen2_vl_attn_forward
+            print("Monkey patch Qwen2VL attention layer")
 
-                patch_vlm_for_ulysses_input_slicing(Qwen2VLModel)
+        if ulysses_sp_size > 1:
+            patch_vlm_for_ulysses_input_slicing(Qwen2VLTextModel)
 
     elif model.config.model_type == "kimi_vl":
         if use_remove_padding or ulysses_sp_size > 1:
@@ -357,43 +321,14 @@ def state_dict(self, *args, **kwargs):
 
         return
 
-    # transformers<=4.47.1
     if use_remove_padding or ulysses_sp_size > 1:
-        if hasattr(module, "_flash_attention_forward"):
+        if hasattr(module, "_flash_attention_forward"):  # transformers <= 4.47.1 or legacy models
             module._flash_attention_forward = _ulysses_flash_attention_forward
             print(f"Monkey patch _flash_attention_forward in {model.__module__}")
         else:
-            if is_transformers_version_in_range(min_version="4.55.0"):
-                from transformers.integrations import flash_attention
-
-                flash_attention._flash_attention_forward = _ulysses_flash_attention_forward_transformers_4_55
-                print(f"Monkey patch _flash_attention_forward in {model.__module__} for new api")
-            else:
-                # 4.48.0 <= transformers <= 4.54.1, Vision attention
-                from transformers.integrations import flash_attention
+            from transformers.integrations import flash_attention
 
-                flash_attention._flash_attention_forward = _ulysses_flash_attention_forward
-                print(f"Monkey patch _flash_attention_forward in {flash_attention.__name__}")
+            flash_attention._flash_attention_forward = _ulysses_flash_attention_forward
+            print(f"Monkey patch _flash_attention_forward in {flash_attention.__name__}")
 
     patch_forward_with_backends(model, use_fused_kernels=use_fused_kernels, fused_kernels_backend=fused_kernels_backend)
-
-
-@lru_cache
-def is_transformers_version_in_range(min_version: Optional[str] = None, max_version: Optional[str] = None) -> bool:
-    try:
-        # Get the installed version of the transformers library
-        transformers_version_str = importlib.metadata.version("transformers")
-    except importlib.metadata.PackageNotFoundError as e:
-        raise ModuleNotFoundError("The `transformers` package is not installed.") from e
-
-    transformers_version = version.parse(transformers_version_str)
-
-    lower_bound_check = True
-    if min_version is not None:
-        lower_bound_check = version.parse(min_version) <= transformers_version
-
-    upper_bound_check = True
-    if max_version is not None:
-        upper_bound_check = transformers_version <= version.parse(max_version)
-
-    return lower_bound_check and upper_bound_check