[megatron] fix: fix qwen2_vl on plain-text data and mix data of plain-text and image-text (volcengine#1999)

### Checklist Before Starting

- [ ] Searched for similar PR(s).
- [ ] Checked PR Title format
  - [ ] In format of: [modules] type: Title
- [ ] modules are in `fsdp, megatron, sglang, vllm, rollout, trainer,
tests, training_utils, recipe, hardware, deployment, ray, worker,
single_controller, misc, perf, model, algo, env, tool, ckpt, doc`
  - [ ] type is in `feat, fix, refactor, chore`
- [ ] can involve multiple modules, seperated by `,` or space, like
`[megatron, fsdp, doc] feat: xxx`

### What does this PR do?


fix qwen2_vl on plain-text data and mix data of plain-text and
image-text, refer to volcengine#1286

### Test


test on gsm8k dataset and mix data of gsm8k and geo3k.

### High-Level Design

> Demonstrate the high-level design if this PR is complex.

### Specific Changes

> List the specific changes.

### API

> Demonstrate how the API changes if any.

### Usage Example

> Provide usage example(s) for easier usage.

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

### Checklist Before Submitting

- [ ] Read the [Contribute
Guide](https://github.com/volcengine/verl?tab=readme-ov-file#contribution-guide).
- [ ] Apply [pre-commit
checks](https://github.com/volcengine/verl?tab=readme-ov-file#code-linting-and-formatting).
- [ ] Add `[BREAKING]` to the PR title `description` if it breaks any
API.
- [ ] Update the documentation about your changes in the
[docs](https://github.com/volcengine/verl/tree/main/docs).
- [ ] New CI unit test(s) are added to cover the code path.
- [ ] Rely on existing unit tests on CI that covers the code path.

Author: MaoChouHJM

verl/models/mcore/model_forward.py

@@ -79,6 +79,8 @@ def gptmodel_forward_qwen2_5_vl(
     assert mpu.get_context_parallel_world_size() == 1, "qwen2_5_vl's context parallel is not accurate yet"
     pre_process = unwrap_model(model).pre_process
     post_process = unwrap_model(model).post_process
+    pixel_values = multi_modal_inputs["pixel_values"].to(input_ids.device) if "pixel_values" in multi_modal_inputs else None
+    image_grid_thw = multi_modal_inputs["image_grid_thw"].to(input_ids.device) if "image_grid_thw" in multi_modal_inputs else None
     if pack_seqs:
         batch_size, seq_len = attention_mask.shape[:2]
         input_ids_rmpad, packed_seq_params = preprocess_packed_seqs(input_ids, attention_mask, pre_process=True)
@@ -88,8 +90,8 @@ def gptmodel_forward_qwen2_5_vl(
             attention_mask=None,
             position_ids=position_ids,
             packed_seq_params=packed_seq_params,
-            pixel_values=multi_modal_inputs["pixel_values"].to(input_ids.device),
-            image_grid_thw=multi_modal_inputs["image_grid_thw"].to(input_ids.device),
+            pixel_values=pixel_values,
+            image_grid_thw=image_grid_thw,
         )
 
         if post_process and logits_processor is not None:
@@ -105,8 +107,8 @@ def gptmodel_forward_qwen2_5_vl(
             input_ids=new_input_ids,
             position_ids=new_position_ids,
             attention_mask=new_attention_mask,
-            pixel_values=multi_modal_inputs["pixel_values"].to(input_ids.device),
-            image_grid_thw=multi_modal_inputs["image_grid_thw"].to(input_ids.device),
+            pixel_values=pixel_values,
+            image_grid_thw=image_grid_thw,
         )
         output = recover_left_padding(output, new_attention_mask, attention_mask, sequence_length, post_process=post_process)
     if value_model and post_process:

verl/workers/actor/megatron_actor.py

@@ -413,7 +413,9 @@ def forward_step(batch_iter, model):
             multi_modal_inputs = {}
             if "multi_modal_inputs" in batch:
                 for key in batch["multi_modal_inputs"][0].keys():
-                    multi_modal_inputs[key] = torch.cat([batch["multi_modal_inputs"][i][key] for i in batch["multi_modal_inputs_idx"]], dim=0)
+                    idxs = batch["multi_modal_inputs_idx"]
+                    mmi = batch["multi_modal_inputs"]
+                    multi_modal_inputs[key] = torch.cat([mmi[idx].get(key) for idx in idxs if mmi[idx].get(key) is not None], dim=0)
             responses = batch["responses"]
             response_length = responses.size(1)
             label = copy.deepcopy(position_ids)