From 36107ec8c9dfda5ee244297bd417afdb46bbbc95 Mon Sep 17 00:00:00 2001 From: alkinun Date: Mon, 18 May 2026 10:40:30 +0300 Subject: [PATCH] Fix ORPO text-only tokenization with processors (#5501) * Fix ORPO text tokenization with processors * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Guard ORPO tokenizer rewrite anchor * Resolve processor pad_token_id and preserve preference data collators for ORPO Two follow-ups so the text-only ORPO + VL processor path works end to end on top of the build_tokenized_answer and tokenize_row rewrites: 1. Add orpo_trainer_processor_pad_token to rewrite processing_class.pad_token_id in ORPOTrainer.__init__ to fall back to processing_class.tokenizer.pad_token_id when the processor itself has no pad_token_id (Qwen3-VL, Gemma-3, etc.). Without this, DPODataCollatorWithPadding(pad_token_id=processing_class.pad_token_id) raises AttributeError before training starts. 2. Stop the outer UnslothORPOTrainer.__init__ collator-swap from clobbering DPODataCollatorWithPadding when the tokenizer is a processor without .pad. The swap to TransformersDataCollatorForLanguageModeling is now only applied to LM-style collators, so ORPO/DPO/CPO/KTO keep their own prompt/chosen/ rejected handling. Otherwise the collator can't pad ORPO rows and raises "You should supply an encoding ... that includes input_ids" at train time. Verified with Qwen3-VL-2B-Instruct ORPO + text-only data (training completes to max_steps, no AttributeError, no collator error) and Llama-3.2-1B-Instruct ORPO (losses and grad-norms bit-exact identical to main, so the change is a true no-op for plain text tokenizers). Extends tests/python/test_orpo_processor_text_tokenizer.py with three new unit tests covering the pad_token_id rewriter. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Wasim Yousef Said Co-authored-by: Daniel Han --- .../test_orpo_processor_text_tokenizer.py | 238 ++++++++++++++++++ unsloth/models/rl.py | 5 +- unsloth/models/rl_replacements.py | 73 ++++++ 3 files changed, 315 insertions(+), 1 deletion(-) create mode 100644 tests/python/test_orpo_processor_text_tokenizer.py diff --git a/tests/python/test_orpo_processor_text_tokenizer.py b/tests/python/test_orpo_processor_text_tokenizer.py new file mode 100644 index 000000000..44c4e26a8 --- /dev/null +++ b/tests/python/test_orpo_processor_text_tokenizer.py @@ -0,0 +1,238 @@ +"""ORPO should use a processor's tokenizer for text-only row tokenization.""" + +import ast +import os +import re + + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +RL_PATH = os.path.join(REPO_ROOT, "unsloth", "models", "rl_replacements.py") + + +def _load_orpo_rewriter(name = "orpo_trainer_text_tokenizer"): + src = open(RL_PATH).read() + tree = ast.parse(src) + ns = {"re": re} + # Materialise sibling module-level assignments (e.g. _PAD_FALLBACK) so + # any rewriter that references them at exec-time can resolve them. + for node in tree.body: + if isinstance(node, ast.Assign): + for target in node.targets: + if isinstance(target, ast.Name) and target.id.startswith("_"): + exec(ast.get_source_segment(src, node), ns) + for node in tree.body: + if isinstance(node, ast.FunctionDef) and node.name == name: + exec(ast.get_source_segment(src, node), ns) + return ns[name] + raise AssertionError(f"{name} not found") + + +class _Tokenizer: + bos_token_id = 1 + eos_token_id = 2 + + def __init__(self): + self.calls = [] + + def __call__(self, text, add_special_tokens = False, **kwargs): + self.calls.append((text, add_special_tokens, kwargs)) + ids = [ord(c) % 31 + 3 for c in text] + return {"input_ids": ids, "attention_mask": [1] * len(ids)} + + +class _Processor: + def __init__(self): + self.tokenizer = _Tokenizer() + + def __call__(self, *args, **kwargs): + raise AssertionError("text-only ORPO tokenization should not call processor") + + +class _Trainer: + def __init__(self): + self.processing_class = _Processor() + self.is_encoder_decoder = False + self.max_length = 2048 + self.max_prompt_length = 1024 + self.max_completion_length = 1024 + self.truncation_mode = "keep_end" + self.label_pad_token_id = -100 + self.padding_value = 0 + + +def _exec_rewritten(function_name, source, extra_ns = None): + rewriter = _load_orpo_rewriter() + rewritten = rewriter(function_name, source) + ns = {} if extra_ns is None else dict(extra_ns) + exec(rewritten, ns) + return ns[function_name] + + +def test_orpo_tokenize_row_returns_original_when_tokenizer_anchor_missing(): + rewriter = _load_orpo_rewriter() + source = """ +def tokenize_row(self, feature, model=None): + output = {} + output["prompt_input_ids"] = self.processing_class(feature["prompt"], add_special_tokens=False)["input_ids"] + return output +""" + + rewritten = rewriter("tokenize_row", source) + + assert rewritten == source + assert "tokenizer(" not in rewritten + + +def test_orpo_build_tokenized_answer_uses_processor_tokenizer(): + source = """ +def build_tokenized_answer(self, prompt, answer): + full_tokenized = self.processing_class(prompt + answer, add_special_tokens=False) + prompt_input_ids = self.processing_class(prompt, add_special_tokens=False)["input_ids"] + return full_tokenized["input_ids"][len(prompt_input_ids):] +""" + fn = _exec_rewritten("build_tokenized_answer", source) + trainer = _Trainer() + + assert fn(trainer, "a", "b") + assert [call[0] for call in trainer.processing_class.tokenizer.calls] == ["ab", "a"] + + +def test_orpo_tokenize_row_uses_processor_tokenizer(): + source = """ +def tokenize_row(self, feature, model=None): + batch = {} + prompt = feature["prompt"] + chosen = feature["chosen"] + rejected = feature["rejected"] + if not self.is_encoder_decoder: + prompt_tokens = self.processing_class(prompt, add_special_tokens=False) + prompt_tokens = {f"prompt_{k}": v for k, v in prompt_tokens.items()} + chosen_tokens = self.build_tokenized_answer(prompt, chosen) + rejected_tokens = self.build_tokenized_answer(prompt, rejected) + prompt_len_input_ids = len(prompt_tokens["prompt_input_ids"]) + chosen_prompt_len_input_ids = len(chosen_tokens["prompt_input_ids"]) + rejected_prompt_len_input_ids = len(rejected_tokens["prompt_input_ids"]) + prompt_tokens, chosen_tokens, rejected_tokens = add_bos_token_if_needed( + self.processing_class.bos_token_id, + prompt_len_input_ids, + prompt_tokens, + chosen_prompt_len_input_ids, + chosen_tokens, + rejected_prompt_len_input_ids, + rejected_tokens, + ) + chosen_tokens, rejected_tokens = add_eos_token_if_needed( + self.processing_class.eos_token_id, chosen_tokens, rejected_tokens + ) + batch["prompt_input_ids"] = prompt_tokens["prompt_input_ids"] + batch["chosen_input_ids"] = chosen_tokens["input_ids"] + batch["rejected_input_ids"] = rejected_tokens["input_ids"] + return batch +""" + + def add_bos_token_if_needed(*args): + return args[2], args[4], args[6] + + def add_eos_token_if_needed(eos_token_id, chosen_tokens, rejected_tokens): + chosen_tokens["input_ids"] = chosen_tokens["input_ids"] + [eos_token_id] + rejected_tokens["input_ids"] = rejected_tokens["input_ids"] + [eos_token_id] + return chosen_tokens, rejected_tokens + + trainer = _Trainer() + trainer.build_tokenized_answer = lambda prompt, answer: { + "prompt_input_ids": trainer.processing_class.tokenizer(prompt)["input_ids"], + "input_ids": trainer.processing_class.tokenizer(answer)["input_ids"], + } + fn = _exec_rewritten( + "tokenize_row", + source, + { + "add_bos_token_if_needed": add_bos_token_if_needed, + "add_eos_token_if_needed": add_eos_token_if_needed, + }, + ) + + output = fn(trainer, {"prompt": "p", "chosen": "c", "rejected": "r"}) + + assert output["chosen_input_ids"][-1] == _Tokenizer.eos_token_id + assert [call[0] for call in trainer.processing_class.tokenizer.calls] == [ + "p", + "p", + "c", + "p", + "r", + ] + + +def test_orpo_init_pad_token_id_falls_back_to_tokenizer(): + rewriter = _load_orpo_rewriter("orpo_trainer_processor_pad_token") + source = """ +def __init__(self, processing_class): + data_collator = DPODataCollatorWithPadding( + pad_token_id=processing_class.pad_token_id, + ) + self.padding_value = processing_class.pad_token_id +""" + + rewritten = rewriter("__init__", source) + + assert "processing_class.pad_token_id" not in rewritten + assert "getattr(processing_class, 'pad_token_id'" in rewritten + + class _Processor: + # No pad_token_id at the processor level; only on the inner tokenizer. + class tokenizer: + pad_token_id = 17 + + captured = {} + + def DPODataCollatorWithPadding(**kwargs): + captured["pad_token_id"] = kwargs["pad_token_id"] + return object() + + ns = {"DPODataCollatorWithPadding": DPODataCollatorWithPadding} + exec(rewritten, ns) + + class _Trainer: + pass + + trainer = _Trainer() + ns["__init__"](trainer, _Processor()) + + assert captured["pad_token_id"] == 17 + assert trainer.padding_value == 17 + + +def test_orpo_init_pad_token_id_uses_processor_when_present(): + rewriter = _load_orpo_rewriter("orpo_trainer_processor_pad_token") + source = """ +def __init__(self, processing_class): + self.padding_value = processing_class.pad_token_id +""" + + rewritten = rewriter("__init__", source) + + class _Tokenizer: + # Inner tokenizer must NOT be consulted when the processor exposes + # pad_token_id itself. + pad_token_id = 999 + + class _Processor: + pad_token_id = 42 + tokenizer = _Tokenizer() + + ns = {} + exec(rewritten, ns) + + class _Trainer: + pass + + trainer = _Trainer() + ns["__init__"](trainer, _Processor()) + assert trainer.padding_value == 42 + + +def test_orpo_init_pad_token_id_noop_on_non_init(): + rewriter = _load_orpo_rewriter("orpo_trainer_processor_pad_token") + source = "def tokenize_row(self):\n return processing_class.pad_token_id\n" + assert rewriter("tokenize_row", source) == source diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 31a498ead..8521b7260 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -1184,6 +1184,9 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"): extra_args += data_collator_check # Also check if .pad exists -> if not, and is VLM, then change it! + # Only swap LM/Seq2Seq collators; leave preference collators + # (DPODataCollatorWithPadding etc.) alone so ORPO/DPO/CPO/KTO keep + # their own prompt/chosen/rejected handling. pad_check = ( "if not isinstance(data_collator, UnslothVisionDataCollator):\n" " if not hasattr(__tokenizer, 'pad') and hasattr(__tokenizer, 'tokenizer'):\n" @@ -1192,7 +1195,7 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"): " __tokenizer.tokenizer,\n" " pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None),\n" " )\n" - " else:\n" + " elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling):\n" " data_collator = TransformersDataCollatorForLanguageModeling(\n" " __tokenizer.tokenizer,\n" " mlm = False,\n" diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 6c70a5cc7..53ef25f95 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -503,6 +503,79 @@ def sft_trainer_compute_loss(function_name, function): RL_FUNCTIONS["sft_trainer"].append(sft_trainer_compute_loss) +# Use the underlying text tokenizer for ORPO row tokenization when a +# multimodal processor is supplied as the processing class. +def orpo_trainer_text_tokenizer(function_name, function): + if function_name == "build_tokenized_answer": + function = re.sub( + r"(?m)^([ \t]*)full_tokenized = self\.processing_class\(prompt \+ answer, add_special_tokens=False\)\n" + r'\1prompt_input_ids = self\.processing_class\(prompt, add_special_tokens=False\)\["input_ids"\]\n', + r'\1tokenizer = getattr(self.processing_class, "tokenizer", self.processing_class)' + "\n" + r"\1full_tokenized = tokenizer(prompt + answer, add_special_tokens=False)" + "\n" + r'\1prompt_input_ids = tokenizer(prompt, add_special_tokens=False)["input_ids"]' + "\n", + function, + count = 1, + ) + return function + + if function_name != "tokenize_row": + return function + + if ( + 'tokenizer = getattr(self.processing_class, "tokenizer", self.processing_class)' + not in function + ): + new_function = re.sub( + r"(?m)^([ \t]*)batch = \{\}\n", + r"\1batch = {}" + "\n" + r'\1tokenizer = getattr(self.processing_class, "tokenizer", self.processing_class)' + "\n", + function, + count = 1, + ) + if new_function == function: + return function + function = new_function + function = function.replace("self.processing_class(", "tokenizer(") + function = function.replace( + "self.processing_class.bos_token_id", "tokenizer.bos_token_id" + ) + function = function.replace( + "self.processing_class.eos_token_id", "tokenizer.eos_token_id" + ) + return function + + +RL_FUNCTIONS["orpo_trainer"].append(orpo_trainer_text_tokenizer) + + +# Resolve `processing_class.pad_token_id` through the underlying tokenizer when +# a multimodal processor is supplied (processors lack `pad_token_id`). Without +# this, ORPOTrainer.__init__ raises AttributeError on +# `DPODataCollatorWithPadding(pad_token_id=processing_class.pad_token_id, ...)` +# and on `self.padding_value = ... else processing_class.pad_token_id`. +_PAD_FALLBACK = ( + "(getattr(processing_class, 'pad_token_id', None) " + "if getattr(processing_class, 'pad_token_id', None) is not None " + "else getattr(getattr(processing_class, 'tokenizer', None), 'pad_token_id', None))" +) + + +def orpo_trainer_processor_pad_token(function_name, function): + if function_name != "__init__": + return function + if "processing_class.pad_token_id" not in function: + return function + return function.replace("processing_class.pad_token_id", _PAD_FALLBACK) + + +RL_FUNCTIONS["orpo_trainer"].append(orpo_trainer_processor_pad_token) + + # Fix bare pop("push_to_hub_token") in compiled SFT/IterativeSFT trainer __init__ # On transformers 5.0+, to_dict() no longer includes push_to_hub_token, so bare pop KeyErrors def sft_trainer_push_to_hub_token(function_name, function):