From 4f9c8321a2136e62fd86fe722a544afd534334a5 Mon Sep 17 00:00:00 2001 From: Datta Nimmaturi Date: Wed, 29 Apr 2026 16:45:34 +0530 Subject: [PATCH] Fix DPO trainer multi process hang (#5199) * Fix DPO trainer multi process hang * Fix datacollator error * further dpo vision changes * cleanup * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Harden DPO vision row processing and source rewrites - dpo_trainer_vision_signature_columns: also match TRL 0.22.x layout (image_sizes followed by ref_chosen_logps), so vision keys are not stripped via remove_unused_columns on the originally-affected version. - dpo_trainer_concatenated_inputs: fall back to inserting after the image_sizes block when no token_type_ids anchor follows it. - Apply the same vision model_kwargs forwarding rewrite to _compute_loss_liger via dpo_trainer_compute_loss_liger so the Liger DPO path does not drop pixel_position_ids/image_position_ids/ mm_token_type_ids when args.use_liger_loss is true. - dpo_trainer_vision_process_row: - guard chosen/rejected EOS append with tokenizer.eos_token_id is not None - use features.get("images") and features.get("prompt") to match the existing get on line 164 and avoid KeyError on rows without those keys - drop the torch.is_tensor gate so list-form pixel_position_ids/ image_position_ids returned without return_tensors are still aliased - skip the loop entry for image_position_ids when it was already promoted to pixel_position_ids, so the output dict no longer carries both keys with identical data - dpo_trainer_data_collator_vision_keys: switch from pad_sequence to trl.trainer.utils.pad with padding_side='left' (matches the DPO collator's prompt left-pad) and padding_value=-1 for *_position_ids keys (sentinel for padded patches), 0 otherwise. Skip the key when not every example carries it. Falls back to pad_sequence if trl.pad is unavailable or the tensor rank is too high. - dpo_trainer_prepare_dataset: keep TRL's writer_batch_size=10 when popping num_proc; removing it defaults to 1000 and reintroduces the vision OOM risk that writer_batch_size=10 was set to avoid. * DPO vision row: keep upstream-facing keys and fix patch padding - dpo_trainer_vision_process_row: no longer aliases image_position_ids to pixel_position_ids. Each upstream-emitted vision key is forwarded under its own name. Gemma4 ForConditionalGeneration.forward accepts image_position_ids directly and renames it to pixel_position_ids only at the vision-tower call site, so aliasing in the row helper hid the kwarg the model actually consumes. - dpo_trainer_vision_process_row: extract pixel_values via "in" membership instead of unconditional indexing. With the missing-images path returning [] to the processor, modern processors no longer emit a pixel_values key, and the previous indexing raised KeyError. - dpo_trainer_data_collator_vision_keys: pick padding_side per key family. *_position_ids tensors are patch-aligned to pixel_values (TRL's DataCollatorForPreference right-pads pixel_values), so pad them right with the -1 sentinel; mm_token_type_ids is token-aligned to prompt_input_ids (left-padded by TRL), so pad it left with 0. * DPO vision: handle multi-image prompts and arbitrary-rank collator pad - dpo_trainer_vision_process_row: when a prompt is missing vision placeholders, insert one placeholder per missing image instead of always inserting a single token. Multi-image rows now satisfy the processor's token-vs-image count check rather than under-inserting and tripping the placeholder/feature mismatch. - dpo_trainer_data_collator_vision_keys: drop the dim()<=2 gate around trl.trainer.utils.pad. trl.pad handles arbitrary rank correctly, while the previous fallback to torch.nn.utils.rnn.pad_sequence raised RuntimeError on rank-3 patch-position tensors with mismatched non-leading dimensions. The pad_sequence path remains as a degraded fallback only when trl.pad is unavailable or raises. * DPO vision row: support scalar images and align prompt-aligned aux ids - dpo_trainer_vision_process_row: type-aware normalization of the features['images'] column instead of a truthiness/len check that raised on single image objects (PIL.Image has no __len__) and on numpy ndarrays (truthiness ambiguous). Lists/tuples count as their length, scalar image objects count as one, None counts as zero, and the original value is forwarded to the processor. - dpo_trainer_vision_process_row: when max_prompt_length truncates prompt_input_ids, also slice token_type_ids and mm_token_type_ids by the same [-max_prompt_length:] suffix. Those keys are 1:1 token aligned to prompt_input_ids (Gemma 4 vision attention keys off mm_token_type_ids per modular_gemma4.py), so leaving them at the original length silently misaligned the multimodal mask. * DPO vision row: stop synthesizing vision-token placeholders Pass features['prompt'] and features['images'] straight to the processor without inserting any extra placeholder tokens. The previous helper used processing_class.image_token, which is the right prompt placeholder for Gemma 4 but the wrong one for Gemma 3 (whose prompt placeholder is boi_token while image_token is the inner expansion target). Synthesizing that token also broke multi-image rows: text ended up with N placeholders while the row helper only forwarded the first image's pixel_values via the standard [0] indexing that mirrors upstream TRL process_row, so token vs image-feature counts diverged. Removing the synthesis matches stock TRL behavior; users provide the correct placeholders for their processor in the prompt. * Add tests for DPO vision row processor passthrough * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: Daniel Han --- .../test_dpo_vision_processor_passthrough.py | 149 ++++++++++ unsloth/models/rl_replacements.py | 272 ++++++++++++++++++ 2 files changed, 421 insertions(+) create mode 100644 tests/python/test_dpo_vision_processor_passthrough.py 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 (