From 22cfcecd08e5fd821e479a99285c0c13c0f9d410 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 24 Feb 2026 07:27:52 +0000 Subject: [PATCH] Forward token_type_ids through GRPO compute_loss to grpo_accumulated_loss In transformers 5.x, Gemma3's create_causal_mask_mapping requires token_type_ids during training to build bidirectional attention masks for image vs text tokens. TRL and the Unsloth data collator already provide token_type_ids in the inputs dict, but compute_loss never extracted it or passed it to grpo_accumulated_loss. This change extracts token_type_ids from inputs and forwards it as a kwarg to both grpo_accumulated_loss call sites (new-style with loss_type and backwards-compat path). Also adds token_type_ids handling to the _get_per_token_logps_and_entropies path for completeness, though that path currently short-circuits via compute_efficient=True. Backwards compatible: token_type_ids defaults to None when not present in inputs, and grpo_accumulated_loss accepts it via **kwargs. --- unsloth/models/rl_replacements.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 8208dc922a..af680edea3 100755 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -670,6 +670,7 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): kwargs.get("pixel_attention_mask", None), kwargs.get("image_sizes", None), ) + token_type_ids = kwargs.get("token_type_ids", None) unwrapped_model = self.accelerator.unwrap_model( model, keep_fp32_wrapper = False @@ -735,6 +736,7 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): pixel_values_chunks = [] image_grid_thw_chunks = [] pixel_attention_mask_chunks = [] + token_type_ids_chunks = [] current_pixel_idx = 0 # TRL 0.23.0 batching logic @@ -766,10 +768,19 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): current_pixel_idx = end_pixel_idx + if token_type_ids is not None: + token_type_ids_chunks.append(token_type_ids[start:end]) + else: + token_type_ids_chunks.append(None) + else: pixel_values_chunks.append(None) image_grid_thw_chunks.append(None) pixel_attention_mask_chunks.append(None) + if token_type_ids is not None: + token_type_ids_chunks.append(token_type_ids[start:end]) + else: + token_type_ids_chunks.append(None) if image_sizes is not None and not isinstance(image_sizes, torch.Tensor): image_sizes_chunks = [[size] for size in image_sizes] @@ -794,6 +805,7 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): image_grid_thw_chunks, pixel_attention_mask_chunks, image_sizes_chunks, + token_type_ids_chunks, ) os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" @@ -805,6 +817,7 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): image_grid_thw_chunk, pixel_attention_mask_chunk, image_sizes_chunk, + token_type_ids_chunk, ) in zipped_inputs: with torch.amp.autocast( device_type = "cuda", dtype = self._autocast_dtype @@ -817,6 +830,7 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): image_grid_thw = image_grid_thw_chunk, pixel_attention_mask = pixel_attention_mask_chunk, image_sizes = image_sizes_chunk, + token_type_ids = token_type_ids_chunk, ).logits completion_input_ids_chunk = input_ids_chunk[ @@ -836,6 +850,7 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): image_grid_thw = image_grid_thw_chunk, pixel_attention_mask = pixel_attention_mask_chunk, image_sizes = image_sizes_chunk, + token_type_ids = token_type_ids_chunk, logits_to_keep = logits_to_keep + 1, ).logits @@ -927,6 +942,7 @@ def grpo_trainer_compute_loss(function_name, function): inputs.get("pixel_attention_mask", None), inputs.get("image_sizes", None), ) + token_type_ids = inputs.get("token_type_ids", None) num_items_in_batch = inputs.get("num_items_in_batch", None) sampling_per_token_logps = inputs.get("sampling_per_token_logps", None) current_gradient_accumulation_steps = self.current_gradient_accumulation_steps @@ -1053,6 +1069,7 @@ def grpo_trainer_compute_loss(function_name, function): logit_scale_multiply = logit_scale_multiply, logit_scale_divide = logit_scale_divide, attention_mask = attention_mask, + token_type_ids = token_type_ids, num_items_in_batch = num_items_in_batch, current_gradient_accumulation_steps = current_gradient_accumulation_steps, num_processes = num_processes, @@ -1075,6 +1092,7 @@ def grpo_trainer_compute_loss(function_name, function): logit_scale_multiply = logit_scale_multiply, logit_scale_divide = logit_scale_divide, attention_mask = attention_mask, + token_type_ids = token_type_ids, ) if "train" in self._metrics: mode = "eval" if self.control.should_evaluate else "train"