Cast multimodal forward_kwargs to compute dtype for bf16/fp16 training

Author: akshan-main

tests/test_grpo_trainer.py

@@ -2287,6 +2287,39 @@ def test_single_reward_model_with_single_processing_class(self):
         assert trainer.reward_processing_classes[0] == single_processing_class
 
 
+def test_forward_kwargs_dtype_casting():
+    forward_kwargs = {
+        "pixel_values": torch.randn(1, 3, 224, 224, dtype=torch.float32),
+        "image_grid_thw": torch.tensor([[1, 14, 14]]),
+    }
+
+    for bf16, fp16, expected_dtype in [
+        (True, False, torch.bfloat16),
+        (False, True, torch.float16),
+        (False, False, None),
+    ]:
+        if bf16:
+            compute_dtype = torch.bfloat16
+        elif fp16:
+            compute_dtype = torch.float16
+        else:
+            compute_dtype = None
+
+        if compute_dtype is not None:
+            result = {
+                k: v.to(compute_dtype) if isinstance(v, torch.Tensor) and torch.is_floating_point(v) else v
+                for k, v in forward_kwargs.items()
+            }
+        else:
+            result = forward_kwargs
+
+        if expected_dtype is not None:
+            assert result["pixel_values"].dtype == expected_dtype
+        else:
+            assert result["pixel_values"].dtype == torch.float32
+        assert result["image_grid_thw"].dtype == torch.int64
+
+
 @pytest.mark.slow
 @require_torch_accelerator
 class TestGRPOTrainerSlow(TrlTestCase):

trl/trainer/grpo_trainer.py

@@ -1609,6 +1609,17 @@ def _generate_and_score_completions(
             prompt_inputs = self.processing_class(images=images, text=prompts_text, padding=True, return_tensors="pt")
             prompt_inputs = super()._prepare_inputs(prompt_inputs)
             forward_kwargs = {k: v for k, v in prompt_inputs.items() if k not in ["input_ids", "attention_mask"]}
+            if self.args.bf16:
+                compute_dtype = torch.bfloat16
+            elif self.args.fp16:
+                compute_dtype = torch.float16
+            else:
+                compute_dtype = None
+            if compute_dtype is not None:
+                forward_kwargs = {
+                    k: v.to(compute_dtype) if isinstance(v, torch.Tensor) and torch.is_floating_point(v) else v
+                    for k, v in forward_kwargs.items()
+                }
         else:
             forward_kwargs = {}