From f106eec5e9d3abc6e45c263299c9a55f3cb0bb58 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 10 Feb 2026 05:14:36 -0800 Subject: [PATCH] Fix Gemma3 4B training on transformers 5.x (token_type_ids) (#4017) * Inject token_type_ids for Gemma3 multimodal training on transformers 5.x In transformers 5.x, create_causal_mask_mapping() raises ValueError when is_training=True and token_type_ids is None. When doing text-only SFT on Gemma3 4B (a multimodal model), the dataset_utils detection for _needs_token_type_ids can miss because: - The model is wrapped in PeftModel, so type(model).__module__ points to peft.peft_model instead of transformers - The processing_class is a tokenizer (not Gemma3Processor), so the fallback MRO check resolves to a module without create_causal_mask_mapping This adds a fallback in _unsloth_pre_compute_loss that injects token_type_ids=zeros when: 1. token_type_ids is not already in inputs 2. The inner model config has model_type "gemma3" 3. The model's module has create_causal_mask_mapping (transformers 5.x) 4. The model is in training mode On transformers 4.x, create_causal_mask_mapping does not exist so this check is inert. Depends on: unslothai/unsloth-zoo#488 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: Daniel Hanchen Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- unsloth/models/_utils.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 2b17c52a54..3657226b1c 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -1754,6 +1754,20 @@ def _unsloth_pre_compute_loss(self, model, inputs, *args, **kwargs): "Using gradient accumulation will be very slightly less accurate.\n" "Read more on gradient accumulation issues here: https://unsloth.ai/blog/gradient" ) + # Gemma3 multimodal models in transformers 5.x require token_type_ids during training. + # For text-only SFT, token_type_ids should be all zeros (no image tokens). + if "token_type_ids" not in inputs and "input_ids" in inputs: + _inner = model + for _attr in ("base_model", "model", "model"): + _inner = getattr(_inner, _attr, _inner) + if getattr(getattr(_inner, "config", None), "model_type", "") in ("gemma3",): + import sys as _sys + + _mod = _sys.modules.get(type(_inner).__module__) + _has_ccm = _mod is not None and hasattr(_mod, "create_causal_mask_mapping") + if _has_ccm and _inner.training: + inputs["token_type_ids"] = torch.zeros_like(inputs["input_ids"]) + outputs = self._old_compute_loss(model, inputs, *args, **kwargs) return outputs