[worker] fix: get all multi_modal_inputs keys with in a microbatch (volcengine#3315)

### What does this PR do?

> Add **concise** overview of what this PR aims to achieve or
accomplish. Reference related GitHub issues and PRs that help with the
review.

Address the first issue in
volcengine#3281 (comment)

More work on top of volcengine#1999

Currently, the code gets the keys from the first row within the
microbatch, This can go wrong if the dataset is a mixture of pure-text
with multi-modal, where the first data in the microbatch is a pure-text
one (no `pixel_values` or `image_grid_thw` exists in the key), and the
microbatch still contains multi-modal data.

This PR fixes this issue by collecting all available keys for
`multi_modal_inputs` within the microbatch, and so that we can
concatenate those multi-modal tensors together without ignoring some of
them under the above situation.

### 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

> For changes that can not be tested by CI (e.g., algorithm
implementation, new model support), validate by experiment(s) and show
results like training curve plots, evaluation results, etc.

### 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

> Demonstrate the high-level design if this PR is complex, and list the
specific changes.

### 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).
- [ ] 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).)

Signed-off-by: Hollow Man <[email protected]>

Author: HollowMan6

verl/utils/model.py

@@ -696,6 +696,47 @@ def get_hf_auto_model_class(hf_config):
     return actor_module_class
 
 
+def extract_multi_modal_inputs(
+    batch_data: list[dict[str, torch.Tensor]],
+    indices: Optional[list[int]] = None,
+) -> dict[str, torch.Tensor | list[torch.Tensor]]:
+    """
+    Extract and process multi-modal inputs from a batch.
+
+    Args:
+        batch_data (list[dict[str, torch.Tensor]]): The batch containing potential multi-modal inputs
+        indices (Optional[list[int]]): If provided, only extract inputs at these indices
+
+    Returns:
+        dict[str, torch.Tensor | list[torch.Tensor]]: Processed multi-modal inputs ready for model consumption
+
+    """
+    multi_modal_inputs = {}
+    multi_modal_inputs_collected = {}
+    has_image_bound = False
+
+    selected_batch_data = batch_data
+    if indices is not None:
+        selected_batch_data = [batch_data[i] for i in indices if i < len(batch_data)]
+
+    for inputs in selected_batch_data:
+        if "image_bound" in inputs:
+            has_image_bound = True
+        for key, value in inputs.items():
+            if value is not None:
+                if key not in multi_modal_inputs_collected:
+                    multi_modal_inputs_collected[key] = []
+                multi_modal_inputs_collected[key].append(value)
+
+    for key, values in multi_modal_inputs_collected.items():
+        if has_image_bound:  # minicpm-o logic
+            multi_modal_inputs[key] = values
+        else:
+            multi_modal_inputs[key] = torch.cat(values, dim=0)
+
+    return multi_modal_inputs
+
+
 @dataclass
 class CausalLMOutputForPPO(CausalLMOutputWithPast):
     log_probs: Optional[torch.FloatTensor] = None

verl/workers/actor/dp_actor.py

@@ -98,14 +98,9 @@ def _forward_micro_batch(
         response_length = micro_batch["responses"].size(-1)
         multi_modal_inputs = {}
         if "multi_modal_inputs" in micro_batch.keys():
-            if "image_bound" in micro_batch["multi_modal_inputs"][0]:  # minicpm-o logic
-                for key in micro_batch["multi_modal_inputs"][0].keys():
-                    multi_modal_inputs[key] = [inputs[key] for inputs in micro_batch["multi_modal_inputs"]]
-            else:
-                for key in micro_batch["multi_modal_inputs"][0].keys():
-                    multi_modal_inputs[key] = torch.cat(
-                        [inputs[key] for inputs in micro_batch["multi_modal_inputs"]], dim=0
-                    )
+            from verl.utils.model import extract_multi_modal_inputs
+
+            multi_modal_inputs = extract_multi_modal_inputs(micro_batch["multi_modal_inputs"])
 
         with torch.autocast(device_type=self.device_name, dtype=torch.bfloat16):
             input_ids = micro_batch["input_ids"]

verl/workers/actor/megatron_actor.py

@@ -466,12 +466,10 @@ def forward_step(batch_iter, model):
 
             multi_modal_inputs = {}
             if "multi_modal_inputs" in batch:
-                for key in batch["multi_modal_inputs"][0].keys():
-                    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
-                    )
+                from verl.utils.model import extract_multi_modal_inputs
+
+                indices = batch.get("multi_modal_inputs_idx", None)
+                multi_modal_inputs = extract_multi_modal_inputs(batch["multi_modal_inputs"], indices)
             responses = batch["responses"]
             response_length = responses.size(1)
             label = position_ids.clone()

verl/workers/critic/dp_critic.py

@@ -58,10 +58,9 @@ def _forward_micro_batch(self, micro_batch):
         response_length = micro_batch["responses"].size(-1)
         multi_modal_inputs = {}
         if "multi_modal_inputs" in micro_batch.keys():
-            for key in micro_batch["multi_modal_inputs"][0].keys():
-                multi_modal_inputs[key] = torch.cat(
-                    [inputs[key] for inputs in micro_batch["multi_modal_inputs"]], dim=0
-                )
+            from verl.utils.model import extract_multi_modal_inputs
+
+            multi_modal_inputs = extract_multi_modal_inputs(micro_batch["multi_modal_inputs"])
 
         with torch.autocast(device_type=self.device_name, dtype=torch.bfloat16):
             input_ids = micro_batch["input_ids"]

verl/workers/engine/fsdp/transformer_impl.py

@@ -644,14 +644,9 @@ def prepare_model_inputs(self, micro_batch: TensorDict):
 
         multi_modal_inputs = {}
         if "multi_modal_inputs" in micro_batch.keys():
-            if "image_bound" in micro_batch["multi_modal_inputs"][0]:  # minicpm-o logic
-                for key in micro_batch["multi_modal_inputs"][0].keys():
-                    multi_modal_inputs[key] = [inputs[key] for inputs in micro_batch["multi_modal_inputs"]]
-            else:
-                for key in micro_batch["multi_modal_inputs"][0].keys():
-                    multi_modal_inputs[key] = torch.cat(
-                        [inputs[key] for inputs in micro_batch["multi_modal_inputs"]], dim=0
-                    )
+            from verl.utils.model import extract_multi_modal_inputs
+
+            multi_modal_inputs = extract_multi_modal_inputs(micro_batch["multi_modal_inputs"])
 
         input_ids = micro_batch["input_ids"]
         attention_mask = micro_batch["attention_mask"]

verl/workers/engine/megatron/transformer_impl.py

@@ -165,19 +165,14 @@ def _build_megatron_module(self):
         return module
 
     def _build_optimizer(self):
-        from verl.utils.megatron.optimizer import (
-            get_megatron_optimizer,
-            init_megatron_optim_config,
-        )
+        from verl.utils.megatron.optimizer import get_megatron_optimizer, init_megatron_optim_config
 
         optim_config_megatron = init_megatron_optim_config(self.optimizer_config)
         optimizer = get_megatron_optimizer(model=self.module, config=optim_config_megatron)
         return optimizer
 
     def _build_lr_scheduler(self):
-        from verl.utils.megatron.optimizer import (
-            get_megatron_optimizer_param_scheduler,
-        )
+        from verl.utils.megatron.optimizer import get_megatron_optimizer_param_scheduler
 
         optimizer_scheduler = get_megatron_optimizer_param_scheduler(
             optimizer=self.optimizer, config=self.optimizer_config
@@ -495,13 +490,11 @@ def prepare_model_inputs(self, batch: TensorDict):
             ]  # mcore patch recompute qwen2vl's pos ids during forward
 
         multi_modal_inputs = {}
-        if "multi_modal_inputs" in batch.keys():
-            for key in batch["multi_modal_inputs"][0].keys():
-                idxs = batch["multi_modal_inputs_idx"]
-                mmi = batch["multi_modal_inputs"]
-                multi_modal_inputs[key] = torch.cat(
-                    [mmi[idx].get(key).to(input_ids.device) for idx in idxs if mmi[idx].get(key) is not None], dim=0
-                )
+        if "multi_modal_inputs" in batch:
+            from verl.utils.model import extract_multi_modal_inputs
+
+            indices = batch.get("multi_modal_inputs_idx", None)
+            multi_modal_inputs = extract_multi_modal_inputs(batch["multi_modal_inputs"], indices)
 
         return {
             "input_ids": input_ids,

verl/workers/reward_model/megatron/reward_model.py

@@ -281,11 +281,10 @@ 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
-                    )
+                from verl.utils.model import extract_multi_modal_inputs
 
+                indices = batch.get("multi_modal_inputs_idx", None)
+                multi_modal_inputs = extract_multi_modal_inputs(batch["multi_modal_inputs"], indices)
             output = forward_fn(
                 model,
                 input_ids,