From c0cf5f257c2e2bb9c1438364afa6e28f1e521286 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 1 Apr 2026 07:28:24 +0000
Subject: [PATCH 1/5] Fix forward compatibility with transformers 5.x

Three issues fixed:

1. Skip exec-based config patching for transformers >= 5.0

Transformers 5.x config classes use @strict, @auto_docstring, and
interval() decorators/annotations that break exec(inspect.getsource(...)).
Those configs already use rope_parameters (the v5 replacement for
rope_scaling), so the patching is not needed. Gated with a version check
so transformers 4.x behavior is unchanged.

2. Slice position_ids to last token in fast_forward_inference

Transformers 5.x generate() accumulates position_ids as
[batch, full_seq_len] across decode steps instead of [batch, 1].
This causes a shape mismatch when indexing cos/sin for rotary
embeddings: cos[position_ids] produces [batch, full_seq_len, head_dim]
but Qn is [batch, n_heads, 1, head_dim]. Fixed by slicing
position_ids[:, -1:] when shape[-1] > 1. Applied to all model files
with fast_forward_inference: llama, qwen3, falcon_h1, gemma2, cohere,
granite. No-op on transformers 4.x since position_ids is already
[batch, 1]. Training path is unaffected.

3. Handle @strict config kwargs for sequence classification

Transformers 5.x @strict config decorator rejects unexpected kwargs
like num_labels, id2label, and max_position_embeddings passed to model
__init__(). Fixed by setting these on the config object directly and
passing config= to from_pretrained. Also added num_labels routing in
FastModel loader to select AutoModelForSequenceClassification.
---
 unsloth/models/_utils.py    | 68 +++++++++++++++++++++----------------
 unsloth/models/cohere.py    |  3 ++
 unsloth/models/falcon_h1.py |  3 ++
 unsloth/models/gemma2.py    |  3 ++
 unsloth/models/granite.py   |  3 ++
 unsloth/models/llama.py     | 13 +++++--
 unsloth/models/loader.py    |  7 +++-
 unsloth/models/qwen3.py     |  3 ++
 unsloth/models/vision.py    |  9 +++++
 9 files changed, 79 insertions(+), 33 deletions(-)

diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index d296ac7e74..dab9adcab6 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -765,43 +765,51 @@ model_architectures = [
     "falcon_h1",
 ]
 
-for model_name in model_architectures:
-    config_filepath = f"transformers.models.{model_name}.configuration_{model_name}"
-    model_filepath = f"transformers.models.{model_name}.modeling_{model_name}"
-    config_filename = f"{model_name.title().replace('_','')}Config"  # qwen3 arch folder is qwen3_moe but config is Qwen3Config. Need to remove underscore(_) for now
-    try:
-        exec(f"from {config_filepath} import {config_filename}", globals())
-    except:
-        continue
+# 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")
 
-    try:
-        config = inspect.getsource(eval(config_filename))
-    except:
-        continue
-    if "RopeParameters" in config:
+if not _skip_config_exec_patch:
+    for model_name in model_architectures:
+        config_filepath = f"transformers.models.{model_name}.configuration_{model_name}"
+        model_filepath = f"transformers.models.{model_name}.modeling_{model_name}"
+        config_filename = f"{model_name.title().replace('_','')}Config"  # qwen3 arch folder is qwen3_moe but config is Qwen3Config. Need to remove underscore(_) for now
         try:
-            exec(f"from {config_filepath} import RopeParameters", globals())
+            exec(f"from {config_filepath} import {config_filename}", globals())
         except:
             continue
 
-    if "rope_scaling" in config:
-        continue
-    config = re.sub(
-        r"(\*\*kwargs)[\s]{0,}\,[\s]{0,}\)[\s]{0,}\:",
-        r"rope_scaling=None,"
-        r"\n        **kwargs):\n"
-        r"\n        self.rope_scaling = rope_scaling\n",
-        config,
-    )
+        try:
+            config = inspect.getsource(eval(config_filename))
+        except:
+            continue
+        if "RopeParameters" in config:
+            try:
+                exec(f"from {config_filepath} import RopeParameters", globals())
+            except:
+                continue
 
-    # Just for Mistral Nemo
-    if model_name == "mistral":
-        if Version(transformers_version) <= Version("4.42.4"):
-            config = patch_mistral_nemo_config(config)
+        if "rope_scaling" in config:
+            continue
+        config = re.sub(
+            r"(\*\*kwargs)[\s]{0,}\,[\s]{0,}\)[\s]{0,}\:",
+            r"rope_scaling=None,"
+            r"\n        **kwargs):\n"
+            r"\n        self.rope_scaling = rope_scaling\n",
+            config,
+        )
 
-    exec(config, globals())
-    exec(f"import {config_filepath}", globals())
-    exec(f"{config_filepath}.{config_filename} = {config_filename}", globals())
+        # Just for Mistral Nemo
+        if model_name == "mistral":
+            if Version(transformers_version) <= Version("4.42.4"):
+                config = patch_mistral_nemo_config(config)
+
+        exec(config, globals())
+        exec(f"import {config_filepath}", globals())
+        exec(f"{config_filepath}.{config_filename} = {config_filename}", globals())
 # =============================================
 
 # =============================================
diff --git a/unsloth/models/cohere.py b/unsloth/models/cohere.py
index 4251f3acd9..294e8d0c7e 100644
--- a/unsloth/models/cohere.py
+++ b/unsloth/models/cohere.py
@@ -357,6 +357,9 @@ def CohereAttention_fast_forward_inference(
     # cos, sin = self.rotary_emb(Vn, seq_len = kv_seq_len)
     # Qn, Kn = inplace_rope_embedding(Qn, Kn, cos, sin, position_ids)
     cos, sin = self.rotary_emb.get_cached(kv_seq_len, Qn.device.index)
+    # Transformers 5.x: position_ids may be [batch, full_seq_len]; slice to last
+    if position_ids.dim() >= 2 and position_ids.shape[-1] > 1:
+        position_ids = position_ids[:, -1:]
     cos = cos[position_ids].unsqueeze(1)
     sin = sin[position_ids].unsqueeze(1)
     h = self.half_head_dim
diff --git a/unsloth/models/falcon_h1.py b/unsloth/models/falcon_h1.py
index 6e3b16b21b..659d27de54 100644
--- a/unsloth/models/falcon_h1.py
+++ b/unsloth/models/falcon_h1.py
@@ -313,6 +313,9 @@ def FalconH1Attention_fast_forward_inference(
     # or else error
     self.rotary_emb.extend_rope_embedding(Vn, seq_len + 2)
     cos, sin = self.rotary_emb.get_cached(kv_seq_len, Qn.device.index)
+    # Transformers 5.x: position_ids may be [batch, full_seq_len]; slice to last
+    if position_ids.dim() >= 2 and position_ids.shape[-1] > 1:
+        position_ids = position_ids[:, -1:]
     cos = cos[position_ids].unsqueeze(1)
     sin = sin[position_ids].unsqueeze(1)
     h = self.half_head_dim
diff --git a/unsloth/models/gemma2.py b/unsloth/models/gemma2.py
index e59b8d5ebd..720c9a7414 100644
--- a/unsloth/models/gemma2.py
+++ b/unsloth/models/gemma2.py
@@ -394,6 +394,9 @@ def Gemma2Attention_fast_forward_inference(
     # cos, sin = self.rotary_emb(Vn, seq_len = kv_seq_len)
     # Qn, Kn = inplace_rope_embedding(Qn, Kn, cos, sin, position_ids)
     cos, sin = self.rotary_emb.get_cached(kv_seq_len, Qn.device.index)
+    # Transformers 5.x: position_ids may be [batch, full_seq_len]; slice to last
+    if position_ids.dim() >= 2 and position_ids.shape[-1] > 1:
+        position_ids = position_ids[:, -1:]
     cos = cos[position_ids].unsqueeze(1)
     sin = sin[position_ids].unsqueeze(1)
     h = self.half_head_dim
diff --git a/unsloth/models/granite.py b/unsloth/models/granite.py
index 79ac41c43f..fea3dc1b36 100644
--- a/unsloth/models/granite.py
+++ b/unsloth/models/granite.py
@@ -355,6 +355,9 @@ def GraniteAttention_fast_forward_inference(
     # cos, sin = self.rotary_emb(Vn, seq_len = kv_seq_len)
     # Qn, Kn = inplace_rope_embedding(Qn, Kn, cos, sin, position_ids)
     cos, sin = position_embeddings
+    # Transformers 5.x: position_ids may be [batch, full_seq_len]; slice to last
+    if position_ids.dim() >= 2 and position_ids.shape[-1] > 1:
+        position_ids = position_ids[:, -1:]
     cos, sin = cos[position_ids], sin[position_ids]
     h = self.half_head_dim
 
diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 93d93e26d6..2f942e259c 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -496,6 +496,10 @@ def LlamaAttention_fast_forward_inference(
     # ensure correct shape
     if position_ids.dim() == 1:
         position_ids = position_ids[:, None]
+    # Transformers 5.x generate() accumulates position_ids as [batch, full_seq_len]
+    # across decode steps. In single-token inference we only need the last position.
+    if position_ids.shape[-1] > 1:
+        position_ids = position_ids[:, -1:]
     position_ids = position_ids.to(Qn.device)
 
     if rotary_seq_len is None:
@@ -2414,14 +2418,19 @@ class FastLlamaModel:
 
         raise_handler = RaiseUninitialized()
         if num_labels is not None:
+            # Transformers 5.x @strict config classes reject unexpected kwargs
+            # like num_labels and max_position_embeddings. Set on the config
+            # object directly and pass config= instead.
+            model_config.num_labels = num_labels
+            if max_position_embeddings is not None:
+                model_config.max_position_embeddings = max_position_embeddings
             model = AutoModelForSequenceClassification.from_pretrained(
                 model_name,
+                config = model_config,
                 device_map = device_map,
                 # torch_dtype             = dtype, # transformers changed torch_dtype to dtype
-                num_labels = num_labels,
                 # quantization_config     = bnb_config,
                 token = token,
-                max_position_embeddings = max_position_embeddings,
                 trust_remote_code = trust_remote_code,
                 attn_implementation = preferred_attn_impl,
                 **kwargs,
diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py
index b54ceaf842..951fe6d228 100644
--- a/unsloth/models/loader.py
+++ b/unsloth/models/loader.py
@@ -1407,8 +1407,13 @@ class FastModel(FastBaseModel):
             architectures = []
         is_vlm = any(x.endswith("ForConditionalGeneration") for x in architectures)
         is_vlm = is_vlm or hasattr(model_config, "vision_config")
+        # If num_labels is set, use AutoModelForSequenceClassification
+        _num_labels = kwargs.get("num_labels", None)
         if auto_model is None:
-            if is_vlm:
+            if _num_labels is not None:
+                from transformers import AutoModelForSequenceClassification
+                auto_model = AutoModelForSequenceClassification
+            elif is_vlm:
                 # Check if the model's auto_map supports the VLM auto class.
                 # Some VL models (e.g. Nemotron-VL) only register AutoModelForCausalLM
                 # in their auto_map, not AutoModelForImageTextToText/AutoModelForVision2Seq.
diff --git a/unsloth/models/qwen3.py b/unsloth/models/qwen3.py
index b93dddb186..3129483be8 100644
--- a/unsloth/models/qwen3.py
+++ b/unsloth/models/qwen3.py
@@ -302,6 +302,9 @@ def Qwen3Attention_fast_forward_inference(
     # or else error
     self.rotary_emb.extend_rope_embedding(Vn, seq_len + 2)
     cos, sin = self.rotary_emb.get_cached(kv_seq_len, Qn.device.index)
+    # Transformers 5.x: position_ids may be [batch, full_seq_len]; slice to last
+    if position_ids.dim() >= 2 and position_ids.shape[-1] > 1:
+        position_ids = position_ids[:, -1:]
     cos = cos[position_ids].unsqueeze(1)
     sin = sin[position_ids].unsqueeze(1)
     h = self.half_head_dim
diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py
index f558aa3f00..7476a158c0 100644
--- a/unsloth/models/vision.py
+++ b/unsloth/models/vision.py
@@ -788,6 +788,15 @@ class FastBaseModel:
         if not fast_inference:
             # Prevent load_in_fp8 from being forwarded into HF internal model loading
             load_in_fp8 = kwargs.pop("load_in_fp8", None)
+            # Transformers 5.x @strict config classes reject unexpected kwargs.
+            # Move config-level attributes onto the config object directly.
+            _num_labels = kwargs.pop("num_labels", None)
+            if _num_labels is not None:
+                model_config.num_labels = _num_labels
+            for _cfg_key in ("id2label", "label2id", "max_position_embeddings"):
+                _cfg_val = kwargs.pop(_cfg_key, None)
+                if _cfg_val is not None:
+                    setattr(model_config, _cfg_key, _cfg_val)
             model = auto_model.from_pretrained(
                 model_name,
                 config = model_config,

From b506fbd86f82e4fd09fb729235315c5b2f9d3188 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 1 Apr 2026 07:28:55 +0000
Subject: [PATCH 2/5] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 unsloth/models/loader.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py
index 951fe6d228..9b7b1b02c5 100644
--- a/unsloth/models/loader.py
+++ b/unsloth/models/loader.py
@@ -1412,6 +1412,7 @@ class FastModel(FastBaseModel):
         if auto_model is None:
             if _num_labels is not None:
                 from transformers import AutoModelForSequenceClassification
+
                 auto_model = AutoModelForSequenceClassification
             elif is_vlm:
                 # Check if the model's auto_map supports the VLM auto class.

From 2290407d996df10a1ccd48bb398b9420d46e45fc Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 1 Apr 2026 08:31:41 +0000
Subject: [PATCH 3/5] Pass token_type_ids and mm_token_type_ids through GRPO
 VLM path

Transformers 5.x requires token_type_ids for some vision models during
training (e.g. Gemma3 Vision calls create_causal_mask_mapping which
raises ValueError if token_type_ids is None during training). Similarly,
Qwen3VL requires mm_token_type_ids for M-RoPE computation.

Extract both from kwargs in _get_per_token_logps_and_entropies, chunk
them alongside other vision tensors, and pass them to the model forward
call via _extra_vision_kwargs dict. This is a no-op when the tensors
are None (transformers 4.x or non-vision models).
---
 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 9f555416d4..8d7fa91b74 100755
--- a/unsloth/models/rl_replacements.py
+++ b/unsloth/models/rl_replacements.py
@@ -714,6 +714,9 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function):
                 kwargs.get("pixel_attention_mask", None),
                 kwargs.get("image_sizes", None),
             )
+            # Transformers 5.x needs token_type_ids/mm_token_type_ids for some vision models
+            token_type_ids = kwargs.get("token_type_ids", None)
+            mm_token_type_ids = kwargs.get("mm_token_type_ids", None)
 
             unwrapped_model = self.accelerator.unwrap_model(
                 model, keep_fp32_wrapper = False
@@ -831,6 +834,10 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function):
             if logit_scale_divide is None:
                 logit_scale_divide = 0
 
+            # Transformers 5.x needs token_type_ids/mm_token_type_ids for some vision models
+            token_type_ids_chunks = chunk_optional(token_type_ids, B)
+            mm_token_type_ids_chunks = chunk_optional(mm_token_type_ids, B)
+
             zipped_inputs = zip(
                 input_ids_chunks,
                 attention_mask_chunks,
@@ -838,6 +845,8 @@ 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,
+                mm_token_type_ids_chunks,
             )
             os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1"
 
@@ -849,7 +858,14 @@ 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,
+                    mm_token_type_ids_chunk,
                 ) in zipped_inputs:
+                    _extra_vision_kwargs = {}
+                    if token_type_ids_chunk is not None:
+                        _extra_vision_kwargs["token_type_ids"] = token_type_ids_chunk
+                    if mm_token_type_ids_chunk is not None:
+                        _extra_vision_kwargs["mm_token_type_ids"] = mm_token_type_ids_chunk
                     with torch.amp.autocast(
                         device_type = "cuda", dtype = self._autocast_dtype
                     ):
@@ -861,6 +877,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,
+                                **_extra_vision_kwargs,
                             ).logits
 
                             completion_input_ids_chunk = input_ids_chunk[
@@ -893,6 +910,7 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function):
                                 pixel_attention_mask = pixel_attention_mask_chunk,
                                 image_sizes = image_sizes_chunk,
                                 logits_to_keep = logits_to_keep + 1,
+                                **_extra_vision_kwargs,
                             ).logits
 
                             logits_chunk = logits_chunk[:, :-1, :]

From d32ee870a70a4bb3ade8658c76f64ed9698bc0a8 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 1 Apr 2026 08:32:56 +0000
Subject: [PATCH 4/5] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 unsloth/models/rl_replacements.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py
index 8d7fa91b74..faba17339b 100755
--- a/unsloth/models/rl_replacements.py
+++ b/unsloth/models/rl_replacements.py
@@ -865,7 +865,9 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function):
                     if token_type_ids_chunk is not None:
                         _extra_vision_kwargs["token_type_ids"] = token_type_ids_chunk
                     if mm_token_type_ids_chunk is not None:
-                        _extra_vision_kwargs["mm_token_type_ids"] = mm_token_type_ids_chunk
+                        _extra_vision_kwargs["mm_token_type_ids"] = (
+                            mm_token_type_ids_chunk
+                        )
                     with torch.amp.autocast(
                         device_type = "cuda", dtype = self._autocast_dtype
                     ):

From c45d5b6efa97b185b3d51ae62453694c38e85407 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 1 Apr 2026 09:13:01 +0000
Subject: [PATCH 5/5] 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          | 14 ++++++--------
 unsloth/models/rl_replacements.py | 18 ++++++++++++++++++
 2 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index dab9adcab6..9d9075affb 100644
--- a/unsloth/models/_utils.py
+++ b/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}"
diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py
index faba17339b..36d8af27cd 100755
--- a/unsloth/models/rl_replacements.py
+++ b/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: