Fix compute_loss token_type_ids propagation and ModernBERT flex_attention

1. In grpo_trainer_compute_loss, extract token_type_ids and
   mm_token_type_ids from inputs dict and pass them to
   grpo_accumulated_loss. Without this, Gemma3 Vision GRPO fails with
   "token_type_ids is required as a model input when training" because
   the model forward never receives it.

2. In _generate_and_score_completions, extend mm_token_type_ids with
   zeros for completion tokens (parallel to existing token_type_ids
   handling). Also save mm_token_type_ids to the output dict. Without
   this, Qwen3VL GRPO fails with get_rope_index shape mismatch because
   mm_token_type_ids covers only the prompt, not prompt+completion.

3. Add "modernbert" to the flex_attention exclusion list in
   prefer_flex_attn_if_supported. ModernBERT with flex_attention hits a
   CUDA illegal memory access in create_block_mask's torch.compile path.
   Falling back to eager attention avoids the crash.

All changes are no-ops on transformers 4.x (token_type_ids/mm_token_type_ids
are None, and modernbert does not exist in the exclusion list check).

unsloth/models/_utils.py

@@ -249,7 +249,7 @@ def prefer_flex_attn_if_supported(model_class, config):
         # NemotronH: hybrid Mamba-2 + Transformer model that does not
         # support flex_attention (raises NotImplementedError from transformers).
         model_type = getattr(config, "model_type", "") if config else ""
-        if model_type in ("gpt_oss", "mllama", "nemotron_h") or str(
+        if model_type in ("gpt_oss", "mllama", "nemotron_h", "modernbert") or str(
             model_type
         ).startswith("gemma3n"):
             return None
@@ -753,6 +753,11 @@ try:
 except:
     from transformers import PretrainedConfig
 
+# transformers 5.x uses class-level annotations + decorators (@strict, @auto_docstring, interval())
+# in config classes, making exec(inspect.getsource(...)) infeasible. Skip config patching for 5.x
+# since those configs already use rope_parameters (renamed from rope_scaling).
+_skip_config_exec_patch = Version(transformers_version) >= Version("5.0.0")
+
 model_architectures = [
     "llama",
     "mistral",
@@ -765,13 +770,6 @@ model_architectures = [
     "falcon_h1",
 ]
 
-# Transformers 5.x uses class-level annotations with @strict, @auto_docstring,
-# and interval() in config classes. exec(inspect.getsource(...)) fails because
-# those symbols are not in scope. Skip the exec-based config patching for 5.x
-# since those configs already use rope_parameters (the v5 replacement for
-# rope_scaling).
-_skip_config_exec_patch = Version(transformers_version) >= Version("5.0.0")
-
 if not _skip_config_exec_patch:
     for model_name in model_architectures:
         config_filepath = f"transformers.models.{model_name}.configuration_{model_name}"

unsloth/models/rl_replacements.py

@@ -542,6 +542,17 @@ def grpo_trainer__generate_and_score_completions(function_name, function):
 
         function = patched
 
+    # Transformers 5.x: Extend mm_token_type_ids for completion tokens (Qwen3VL M-RoPE)
+    # TRL handles token_type_ids but not mm_token_type_ids
+    _tt_search = 'if "token_type_ids" in forward_kwargs:\n            token_type_ids = forward_kwargs["token_type_ids"]\n            forward_kwargs["token_type_ids"] = torch.cat(\n                [token_type_ids, token_type_ids.new_zeros(completion_ids.shape)], dim=1\n            )'
+    _tt_replace = _tt_search + '\n        if "mm_token_type_ids" in forward_kwargs:\n            mm_tti = forward_kwargs["mm_token_type_ids"]\n            forward_kwargs["mm_token_type_ids"] = torch.cat(\n                [mm_tti, mm_tti.new_zeros(completion_ids.shape)], dim=1\n            )'
+    function = function.replace(_tt_search, _tt_replace)
+
+    # Save mm_token_type_ids to output dict alongside token_type_ids
+    _save_search = 'if "token_type_ids" in forward_kwargs:\n            output["token_type_ids"] = forward_kwargs["token_type_ids"]'
+    _save_replace = _save_search + '\n        if "mm_token_type_ids" in forward_kwargs:\n            output["mm_token_type_ids"] = forward_kwargs["mm_token_type_ids"]'
+    function = function.replace(_save_search, _save_replace)
+
     return function
 
 
@@ -1013,6 +1024,9 @@ def grpo_trainer_compute_loss(function_name, function):
             inputs.get("pixel_attention_mask", None),
             inputs.get("image_sizes", None),
         )
+        # Transformers 5.x needs token_type_ids/mm_token_type_ids for some vision models
+        token_type_ids = inputs.get("token_type_ids", None)
+        mm_token_type_ids = inputs.get("mm_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
@@ -1156,6 +1170,8 @@ def grpo_trainer_compute_loss(function_name, function):
                     current_gradient_accumulation_steps = current_gradient_accumulation_steps,
                     num_processes = num_processes,
                     sampling_per_token_logps = sampling_per_token_logps,
+                    token_type_ids = token_type_ids,
+                    mm_token_type_ids = mm_token_type_ids,
                 )
             else:
                 # to ensure backwards compatibility with trl 0.15.2 and maybe even 0.17
@@ -1174,6 +1190,8 @@ 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,
+                        mm_token_type_ids = mm_token_type_ids,
                     )
                 )
         if "train" in self._metrics: