diff --git a/tests/python/test_dpo_vision_processor_passthrough.py b/tests/python/test_dpo_vision_processor_passthrough.py new file mode 100644 index 0000000000..a4f2e2e12a --- /dev/null +++ b/tests/python/test_dpo_vision_processor_passthrough.py @@ -0,0 +1,149 @@ +"""Verify dpo_trainer_vision_process_row forwards prompt and images verbatim.""" + +import ast +import os + +import numpy as np + + +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_helpers(): + src = open(RL_PATH).read() + tree = ast.parse(src) + import torch as _torch + + ns = {"torch": _torch} + for node in tree.body: + if isinstance(node, ast.Assign) and any( + isinstance(t, ast.Name) and t.id == "_DPO_VISION_KEYS" for t in node.targets + ): + exec(ast.get_source_segment(src, node), ns) + for node in tree.body: + if isinstance(node, ast.FunctionDef) and node.name.startswith( + ("dpo_trainer_", "_dpo_trainer_") + ): + exec(ast.get_source_segment(src, node), ns) + return ns + + +class _Tok: + eos_token_id = 99 + bos_token_id = None + + def __call__(self, t, add_special_tokens = False): + return {"input_ids": [10]} + + +class _Capture: + image_token = "" + boi_token = "" + + def __init__(self): + self.tokenizer = _Tok() + self.last_text = None + self.last_images = "__sentinel__" + + def __call__(self, images = None, text = None, add_special_tokens = False): + self.last_text = text + self.last_images = images + out = {"input_ids": [[1, 2]]} + if images is not None: + out["pixel_values"] = [object()] + return out + + +def test_prompt_passes_through_without_image_token_synthesis(): + ns = _load_helpers() + proc = _Capture() + ns["dpo_trainer_vision_process_row"]( + {"prompt": "describe", "chosen": "c", "rejected": "r", "images": ["i"]}, + proc, + ) + assert proc.last_text == "describe" + + +def test_prompt_with_existing_image_token_unchanged(): + ns = _load_helpers() + proc = _Capture() + ns["dpo_trainer_vision_process_row"]( + {"prompt": " describe", "chosen": "c", "rejected": "r", "images": ["i"]}, + proc, + ) + assert proc.last_text == " describe" + + +def test_gemma3_style_boi_token_prompt_not_corrupted(): + ns = _load_helpers() + proc = _Capture() + ns["dpo_trainer_vision_process_row"]( + {"prompt": " describe", "chosen": "c", "rejected": "r", "images": ["i"]}, + proc, + ) + assert proc.last_text == " describe" + assert "" not in proc.last_text + + +def test_multi_image_prompt_unchanged_no_extra_placeholders(): + ns = _load_helpers() + proc = _Capture() + ns["dpo_trainer_vision_process_row"]( + { + "prompt": "compare", + "chosen": "c", + "rejected": "r", + "images": ["a", "b", "c"], + }, + proc, + ) + assert proc.last_text == "compare" + + +def test_list_images_forwarded_verbatim(): + ns = _load_helpers() + proc = _Capture() + payload = ["a", "b"] + ns["dpo_trainer_vision_process_row"]( + {"prompt": "p", "chosen": "c", "rejected": "r", "images": payload}, + proc, + ) + assert proc.last_images is payload + + +def test_single_pil_like_image_forwarded_verbatim(): + ns = _load_helpers() + + class PIL: + def __bool__(self): + return True + + proc = _Capture() + pil = PIL() + ns["dpo_trainer_vision_process_row"]( + {"prompt": "p", "chosen": "c", "rejected": "r", "images": pil}, + proc, + ) + assert proc.last_images is pil + + +def test_numpy_ndarray_image_forwarded_verbatim(): + ns = _load_helpers() + proc = _Capture() + arr = np.zeros((2, 3, 3), dtype = np.uint8) + ns["dpo_trainer_vision_process_row"]( + {"prompt": "p", "chosen": "c", "rejected": "r", "images": arr}, + proc, + ) + assert proc.last_images is arr + + +def test_missing_images_key_passes_none_to_processor(): + ns = _load_helpers() + proc = _Capture() + ns["dpo_trainer_vision_process_row"]( + {"prompt": "p", "chosen": "c", "rejected": "r"}, + proc, + ) + assert proc.last_images is None diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 4d36af62cc..5d2c4151cf 100755 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -55,6 +55,12 @@ RL_CONFIG_CHANGES = defaultdict(list) RL_METRICS_CHANGES = defaultdict(list) RL_ADDITIONAL_FUNCTIONS = defaultdict(list) +_DPO_VISION_KEYS = ( + "pixel_position_ids", + "image_position_ids", + "mm_token_type_ids", +) + torch_compile_options = { "epilogue_fusion": True, "max_autotune": False, # I saw speedups, but not sure if this has issues in collab @@ -120,6 +126,272 @@ def dpo_trainer_fix_columns(call_args, extra_args): RL_EXTRA_ARGS["dpo_trainer"].append(dpo_trainer_fix_columns) +def dpo_trainer_fix_data_collator(call_args, extra_args): + if ( + "data_collator" in call_args + and "train_dataset" in call_args + and "processing_class" in call_args + ): + fix_collator = ( + "if hasattr(train_dataset, 'column_names'):\n" + " column_names = set(train_dataset.column_names)\n" + " is_dpo_dataset = ({'chosen', 'rejected'}.issubset(column_names) or\n" + " {'prompt_input_ids', 'chosen_input_ids', 'rejected_input_ids'}.issubset(column_names))\n" + " if is_dpo_dataset and isinstance(data_collator, TransformersDataCollatorForLanguageModeling):\n" + " data_collator = None\n" + " del is_dpo_dataset, column_names\n" + ) + return fix_collator + return "" + + +RL_EXTRA_ARGS["dpo_trainer"].append(dpo_trainer_fix_data_collator) + + +def dpo_trainer_vision_process_row( + features, + processing_class, + max_prompt_length = None, + max_completion_length = None, + add_special_tokens = True, + is_chat = False, +): + text = features.get("prompt", "") + images = features.get("images") + processor, tokenizer = processing_class, processing_class.tokenizer + processed_features = processor( + images = images, + text = text, + add_special_tokens = False, + ) + + prompt_input_ids = processed_features["input_ids"][0] + chosen_input_ids = tokenizer(features["chosen"], add_special_tokens = False)[ + "input_ids" + ] + rejected_input_ids = tokenizer(features["rejected"], add_special_tokens = False)[ + "input_ids" + ] + + if add_special_tokens: + if tokenizer.bos_token_id is not None: + prompt_input_ids = [tokenizer.bos_token_id] + prompt_input_ids + if tokenizer.eos_token_id is not None: + prompt_input_ids = prompt_input_ids + [tokenizer.eos_token_id] + if not is_chat and tokenizer.eos_token_id is not None: + chosen_input_ids = chosen_input_ids + [tokenizer.eos_token_id] + rejected_input_ids = rejected_input_ids + [tokenizer.eos_token_id] + + if max_prompt_length is not None: + prompt_input_ids = prompt_input_ids[-max_prompt_length:] + if max_completion_length is not None: + chosen_input_ids = chosen_input_ids[:max_completion_length] + rejected_input_ids = rejected_input_ids[:max_completion_length] + + output = { + "prompt_input_ids": prompt_input_ids, + "chosen_input_ids": chosen_input_ids, + "rejected_input_ids": rejected_input_ids, + } + if "pixel_values" in processed_features: + output["pixel_values"] = processed_features["pixel_values"][0] + if "pixel_attention_mask" in processed_features: + output["pixel_attention_mask"] = processed_features["pixel_attention_mask"][0] + if "image_sizes" in processed_features: + output["image_sizes"] = processed_features["image_sizes"][0] + if "token_type_ids" in processed_features: + token_type_ids = processed_features["token_type_ids"][0] + if max_prompt_length is not None: + token_type_ids = token_type_ids[-max_prompt_length:] + output["token_type_ids"] = token_type_ids + if "pixel_position_ids" in processed_features: + output["pixel_position_ids"] = processed_features["pixel_position_ids"][0] + if "image_position_ids" in processed_features: + output["image_position_ids"] = processed_features["image_position_ids"][0] + if "mm_token_type_ids" in processed_features: + mm_token_type_ids = processed_features["mm_token_type_ids"][0] + if max_prompt_length is not None: + mm_token_type_ids = mm_token_type_ids[-max_prompt_length:] + output["mm_token_type_ids"] = mm_token_type_ids + + return output + + +def dpo_trainer_vision_signature_columns(function_name, function): + if function_name != "_set_signature_columns_if_needed": + return function + + if all(_k in function for _k in _DPO_VISION_KEYS): + return function + + _extra_columns = "".join(f' "{_k}",\n' for _k in _DPO_VISION_KEYS) + new_function = function.replace( + ' "image_sizes",\n' ' "token_type_ids",\n', + f' "image_sizes",\n' + f"{_extra_columns}" + f' "token_type_ids",\n', + ) + if new_function != function: + return new_function + return function.replace( + ' "image_sizes",\n' ' "ref_chosen_logps",\n', + f' "image_sizes",\n' + f"{_extra_columns}" + f' "ref_chosen_logps",\n', + ) + + +def dpo_trainer_concatenated_inputs(function_name, function): + if function_name != "concatenated_inputs": + return function + + if all(_k in function for _k in _DPO_VISION_KEYS): + return function + + _extra_inputs = "".join( + f' if "{_k}" in batch:\n' + f' output["{_k}"] = torch.cat((batch["{_k}"], batch["{_k}"]), dim=0)\n' + for _k in _DPO_VISION_KEYS + ) + + image_sizes_block = ( + ' if "image_sizes" in batch:\n' + ' output["image_sizes"] = torch.cat([batch["image_sizes"], batch["image_sizes"]], dim=0)\n' + ) + new_function = function.replace( + image_sizes_block + ' if "token_type_ids" in batch:\n', + image_sizes_block + _extra_inputs + ' if "token_type_ids" in batch:\n', + ) + if new_function != function: + return new_function + if image_sizes_block in function: + return function.replace(image_sizes_block, image_sizes_block + _extra_inputs, 1) + return function + + +def _dpo_trainer_extend_vision_model_kwargs(function): + if all(_k in function for _k in _DPO_VISION_KEYS): + return function + + _extra_forward = "".join( + f' if "{_k}" in concatenated_batch:\n' + f' model_kwargs["{_k}"] = concatenated_batch["{_k}"]\n' + for _k in ( + "pixel_values", + "pixel_attention_mask", + "image_sizes", + *_DPO_VISION_KEYS, + ) + ) + + return function.replace( + ' if "pixel_values" in concatenated_batch:\n' + ' model_kwargs["pixel_values"] = concatenated_batch["pixel_values"]\n' + ' if "pixel_attention_mask" in concatenated_batch:\n' + ' model_kwargs["pixel_attention_mask"] = concatenated_batch["pixel_attention_mask"]\n' + ' if "image_sizes" in concatenated_batch:\n' + ' model_kwargs["image_sizes"] = concatenated_batch["image_sizes"]\n', + f"{_extra_forward}", + ) + + +def dpo_trainer_concatenated_forward(function_name, function): + if function_name != "concatenated_forward": + return function + return _dpo_trainer_extend_vision_model_kwargs(function) + + +def dpo_trainer_compute_loss_liger(function_name, function): + if function_name != "_compute_loss_liger": + return function + return _dpo_trainer_extend_vision_model_kwargs(function) + + +def dpo_trainer_data_collator_vision_keys(call_args, extra_args): + if "data_collator" not in call_args: + return "" + + _vision_keys = str(_DPO_VISION_KEYS) + return ( + "from trl.trainer.dpo_trainer import DataCollatorForPreference\n" + "if not hasattr(DataCollatorForPreference, '_unsloth_vision_keys_patch'):\n" + " _old_dpo_collator_torch_call = DataCollatorForPreference.torch_call\n" + "\n" + " def _unsloth_dpo_torch_call(self, examples):\n" + " output = _old_dpo_collator_torch_call(self, examples)\n" + " import torch as _unsloth_torch\n" + " try:\n" + " from trl.trainer.utils import pad as _unsloth_trl_pad\n" + " except Exception:\n" + " _unsloth_trl_pad = None\n" + " for _k in " + _vision_keys + ":\n" + " if not all(_k in example for example in examples):\n" + " continue\n" + " _is_position_key = _k.endswith('position_ids')\n" + " _padding_value = -1 if _is_position_key else 0\n" + " _padding_side = 'right' if _is_position_key else 'left'\n" + " _values = [_unsloth_torch.as_tensor(example[_k]) for example in examples]\n" + " try:\n" + " if _unsloth_trl_pad is not None:\n" + " output[_k] = _unsloth_trl_pad(_values, padding_value=_padding_value, padding_side=_padding_side)\n" + " else:\n" + " from torch.nn.utils.rnn import pad_sequence as _unsloth_pad_sequence\n" + " output[_k] = _unsloth_pad_sequence(_values, batch_first=True, padding_value=_padding_value)\n" + " except Exception:\n" + " from torch.nn.utils.rnn import pad_sequence as _unsloth_pad_sequence\n" + " output[_k] = _unsloth_pad_sequence(_values, batch_first=True, padding_value=_padding_value)\n" + " return output\n" + "\n" + " DataCollatorForPreference.torch_call = _unsloth_dpo_torch_call\n" + " DataCollatorForPreference._unsloth_vision_keys_patch = True\n" + ) + + +def dpo_trainer_prepare_dataset(function_name, function): + if function_name != "_prepare_dataset": + return function + + legacy_call = "self.tokenize_row if not self.is_vision_model else self.process_row" + if legacy_call not in function: + return function + + function = function.replace( + legacy_call, + "self.tokenize_row if not self.is_vision_model else dpo_trainer_vision_process_row", + ) + + legacy_tokenize_block = ( + " # Tokenize the dataset\n" + " if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc`\n" + ' map_kwargs["desc"] = f"Tokenizing {dataset_name} dataset"\n' + "\n" + " dataset = dataset.map(\n" + " self.tokenize_row if not self.is_vision_model else dpo_trainer_vision_process_row,\n" + ) + patched_tokenize_block = ( + " # Tokenize the dataset\n" + " if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc`\n" + ' map_kwargs["desc"] = f"Tokenizing {dataset_name} dataset"\n' + " if self.is_vision_model:\n" + ' map_kwargs.pop("num_proc", None)\n' + "\n" + " dataset = dataset.map(\n" + " self.tokenize_row if not self.is_vision_model else dpo_trainer_vision_process_row,\n" + ) + if legacy_tokenize_block in function: + function = function.replace(legacy_tokenize_block, patched_tokenize_block, 1) + return function + + +RL_FUNCTIONS["dpo_trainer"].append(dpo_trainer_prepare_dataset) +RL_PRE_ITEMS["dpo_trainer"].append(inspect.getsource(dpo_trainer_vision_process_row)) +RL_FUNCTIONS["dpo_trainer"].append(dpo_trainer_vision_signature_columns) +RL_FUNCTIONS["dpo_trainer"].append(dpo_trainer_concatenated_inputs) +RL_FUNCTIONS["dpo_trainer"].append(dpo_trainer_concatenated_forward) +RL_FUNCTIONS["dpo_trainer"].append(dpo_trainer_compute_loss_liger) +RL_EXTRA_ARGS["dpo_trainer"].append(dpo_trainer_data_collator_vision_keys) + + # Fix tokenizer double BOS def sft_trainer_prepare_dataset(function_name, function): if (