[data] better mm data collate (#424)

Author: hiyouga

verl/workers/actor/dp_actor.py

@@ -26,7 +26,7 @@
 from torch import nn
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 
-from ...protocol import DataProto
+from ...protocol import DataProto, batch_collate
 from ...trainer.core_algos import average_loss, compute_kl, compute_policy_loss
 from ...utils import torch_functional as VF
 from ...utils.py_functional import append_to_dict
@@ -81,15 +81,10 @@ def _forward_micro_batch(self, micro_batch: Dict[str, torch.Tensor], temperature
 
         multi_modal_inputs = defaultdict(list)
         if "multi_modal_inputs" in micro_batch:
-            for input_dict in micro_batch["multi_modal_inputs"]:
-                for key, value in input_dict.items():
-                    multi_modal_inputs[key].append(value)
-
-            for key, value in multi_modal_inputs.items():
-                if len(value) != 0:
-                    multi_modal_inputs[key] = torch.cat(value, dim=0)
-                else:
-                    multi_modal_inputs[key] = None
+            multi_modal_inputs = batch_collate(micro_batch["multi_modal_inputs"])
+            multi_modal_inputs = {key: torch.cat(value, dim=0) for key, value in multi_modal_inputs.items()}
+        else:
+            multi_modal_inputs = {}
 
         if self.config.padding_free:
             input_ids_rmpad, indices, *_ = unpad_input(input_ids.unsqueeze(-1), attention_mask)  # (total_nnz, 1)

verl/workers/critic/dp_critic.py

@@ -25,7 +25,7 @@
 from torch import nn
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 
-from ...protocol import DataProto
+from ...protocol import DataProto, batch_collate
 from ...trainer.core_algos import compute_value_loss
 from ...utils.py_functional import append_to_dict
 from ...utils.seqlen_balancing import prepare_dynamic_batch, restore_dynamic_batch
@@ -61,17 +61,11 @@ def _forward_micro_batch(self, micro_batch: Dict[str, torch.Tensor]) -> torch.Te
         if position_ids.dim() == 3:  # qwen2vl mrope
             position_ids = position_ids.transpose(0, 1)  # (bsz, 3, seqlen) -> (3, bsz, seqlen)
 
-        multi_modal_inputs = defaultdict(list)
         if "multi_modal_inputs" in micro_batch:
-            for input_dict in micro_batch["multi_modal_inputs"]:
-                for key, value in input_dict.items():
-                    multi_modal_inputs[key].append(value)
-
-            for key, value in multi_modal_inputs.items():
-                if len(value) != 0:
-                    multi_modal_inputs[key] = torch.cat(value, dim=0)
-                else:
-                    multi_modal_inputs[key] = None
+            multi_modal_inputs = batch_collate(micro_batch["multi_modal_inputs"])
+            multi_modal_inputs = {key: torch.cat(value, dim=0) for key, value in multi_modal_inputs.items()}
+        else:
+            multi_modal_inputs = {}
 
         if self.config.padding_free:
             input_ids_rmpad, indices, *_ = unpad_input(