ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix ORPO text-only tokenization with processors (#5501)

Browse source 
* 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 <wasimysdev@gmail.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
This commit is contained in:
alkinun 2026-05-18 10:40:30 +03:00 • committed by GitHub
commit 36107ec8c9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 315 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

238
tests/python/test_orpo_processor_text_tokenizer.py Normal file
View file

@ -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

5
unsloth/models/rl.py
View file

@ -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"

73
unsloth/models/rl_replacements.py
View file

@ -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):