From a1c37721c3e4b9d0a2546b4bc47f9f5e85cb069d Mon Sep 17 00:00:00 2001 From: Strongich Date: Sun, 5 Apr 2026 15:28:07 +0200 Subject: [PATCH 01/39] Add VLM support to GOLDTrainer (without vLLM yet) --- tests/experimental/test_gold_trainer.py | 392 +++++++++++++++++++++++- trl/experimental/gold/gold_trainer.py | 88 +++++- trl/experimental/utils.py | 161 +++++++++- 3 files changed, 632 insertions(+), 9 deletions(-) diff --git a/tests/experimental/test_gold_trainer.py b/tests/experimental/test_gold_trainer.py index d7e32056323..3257ad2f99f 100644 --- a/tests/experimental/test_gold_trainer.py +++ b/tests/experimental/test_gold_trainer.py @@ -16,12 +16,12 @@ import pytest import torch -from datasets import load_dataset -from transformers import AutoTokenizer +from datasets import Dataset, load_dataset +from transformers import AutoProcessor, AutoTokenizer from trl.experimental.gold import gold_trainer as gold_trainer_module from trl.experimental.gold.gold_trainer import GOLDTrainer, ULDLoss, build_teacher_inputs_from_texts -from trl.experimental.utils import DataCollatorForChatML +from trl.experimental.utils import DataCollatorForChatML, DataCollatorForVisionLanguageChatML @pytest.fixture(scope="module") @@ -271,6 +271,35 @@ def smollm_tokenizer(): return tokenizer +@pytest.fixture(scope="session") +def smolvlm_processor(): + processor = AutoProcessor.from_pretrained("HuggingFaceTB/SmolVLM-256M-Instruct") + if processor.tokenizer.pad_token is None: + processor.tokenizer.pad_token = processor.tokenizer.eos_token + return processor + + +@pytest.fixture(scope="session") +def qwen3_vl_processor(): + processor = AutoProcessor.from_pretrained("Qwen/Qwen3-VL-2B-Instruct") + if processor.tokenizer.pad_token is None: + processor.tokenizer.pad_token = processor.tokenizer.eos_token + return processor + + +@pytest.fixture(scope="module") +def vlm_examples(): + try: + dataset = load_dataset( + "trl-internal-testing/zen-image", + "conversational_prompt_completion", + split="train[:3]", + ) + except Exception as exc: # pragma: no cover - network/environment dependent + pytest.skip(f"zen-image dataset unavailable: {exc}") + return [dict(row) for row in dataset] + + def encode_prompt_completion(tokenizer, prompt, completion): prompt_ids = tokenizer(prompt, add_special_tokens=False)["input_ids"] completion_ids = tokenizer(completion, add_special_tokens=False)["input_ids"] @@ -575,12 +604,17 @@ def __init__(self, **kwargs): vllm_sync_frequency=1, ) + # A minimal dataset is required because GOLDTrainer inspects the first example at init + # to detect whether the dataset contains images for VLM, so None dosn't pass + dummy_dataset = Dataset.from_dict({"messages": [["dummy"]]}) + teacher_model = DummyTeacherModel() GOLDTrainer( model=DummyStudentModel(), teacher_model=teacher_model, args=args, data_collator=object(), + train_dataset=dummy_dataset, processing_class=DummyProcessingClass(), ) @@ -942,3 +976,355 @@ def test_uldloss_hybrid_config_beta_zero(llama_tokenizer, qwen_tokenizer): expected = config.uld_hybrid_unmatched_weight * loss_fn.last_unmatched_loss torch.testing.assert_close(loss, expected, atol=1e-6, rtol=1e-5) + + +# ────────────────────────────────────────────────────────────────────────────── +# VLM tests +# ────────────────────────────────────────────────────────────────────────────── + + +def test_vlm_alignment_groups_cover_all_tokens_smolvlm_qwen3vl(smolvlm_processor, qwen3_vl_processor, vlm_examples): + student_tokenizer = smolvlm_processor.tokenizer + teacher_tokenizer = qwen3_vl_processor.tokenizer + + collator = DataCollatorForVisionLanguageChatML(processor=smolvlm_processor, max_length=2048) + batch = collator(vlm_examples) + + config = build_config() + loss = ULDLoss(config, student_tokenizer=student_tokenizer, teacher_tokenizer=teacher_tokenizer) + + teacher_input_ids, teacher_labels, _ = _teacher_inputs_from_collator(student_tokenizer, teacher_tokenizer, batch) + + _assert_alignment_covers_completion(loss, batch, teacher_input_ids, teacher_labels) + + +def test_gold_trainer_init_rejects_llm_with_vision_dataset(monkeypatch): + """GOLDTrainer should raise ValueError when a text-only model receives a vision dataset.""" + + class DummyStudentModel: + def __init__(self): + self.config = SimpleNamespace(_name_or_path="student", vocab_size=17) + self.generation_config = SimpleNamespace(eos_token_id=2) + self.name_or_path = "student" + + class DummyTeacherModel: + def __init__(self): + self.resized_to = None + + def resize_token_embeddings(self, vocab_size): + self.resized_to = vocab_size + + def fake_sft_init( + self, + model, + args=None, + data_collator=None, + train_dataset=None, + eval_dataset=None, + processing_class=None, + compute_metrics=None, + callbacks=None, + optimizers=None, + preprocess_logits_for_metrics=None, + peft_config=None, + ): + del data_collator, train_dataset, eval_dataset, compute_metrics, callbacks, optimizers + del preprocess_logits_for_metrics, peft_config + self.model = model + self.args = args + self.processing_class = processing_class + self.accelerator = SimpleNamespace( + device=torch.device("cpu"), + num_processes=1, + prepare_model=lambda module, evaluation_mode=True: module, + ) + self.is_deepspeed_enabled = False + self.is_fsdp_enabled = False + + monkeypatch.setattr(gold_trainer_module.SFTTrainer, "__init__", fake_sft_init) + + tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM3-3B") + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + + # Dataset with an "image" key triggers vision detection + vision_dataset = Dataset.from_dict({"messages": [["dummy"]], "image": ["fake_image"]}) + + args = SimpleNamespace( + model_init_kwargs=None, + max_length=128, + use_liger_kernel=False, + teacher_model_init_kwargs=None, + use_uld_loss=False, + teacher_tokenizer_name_or_path=None, + teacher_model_revision=None, + disable_dropout=False, + lmbda=1.0, + beta=0.5, + temperature=1.0, + top_p=1.0, + seq_kd=False, + num_generations=1, + use_transformers_paged=False, + max_completion_length=16, + top_k=0, + log_completions=False, + log_completions_steps=100, + wandb_log_unique_prompts=True, + num_completions_to_print=None, + per_device_train_batch_size=1, + gradient_accumulation_steps=1, + use_vllm=False, + ) + + with pytest.raises(ValueError, match="vision-related"): + GOLDTrainer( + model=DummyStudentModel(), + teacher_model=DummyTeacherModel(), + args=args, + train_dataset=vision_dataset, + processing_class=tokenizer, + ) + + +def _get_assistant_texts(examples): + """Extract assistant text content from examples, handling both plain string and multimodal format.""" + texts = [] + for example in examples: + content = example["completion"][-1]["content"] + if isinstance(content, list): + texts.append("".join(part["text"] for part in content if "text" in part)) + else: + texts.append(content) + return texts + + +def test_vlm_chatml_collator_preserves_completion_smolvlm(smolvlm_processor, qwen3_vl_processor, vlm_examples): + # 2048 to not truncate the completion tokens + collator = DataCollatorForVisionLanguageChatML(processor=smolvlm_processor, max_length=2048) + batch = collator(vlm_examples) + + # Verify basic batch structure + assert "input_ids" in batch + assert "labels" in batch + assert "prompts" in batch + assert "prompt_attention_mask" in batch + assert "pixel_values" in batch + assert "original_prompt_text" in batch + assert "original_completion_text" in batch + + # Verify completions are preserved in decoded output + assistant_texts = _get_assistant_texts(vlm_examples) + decoded_batch = smolvlm_processor.tokenizer.batch_decode(batch["input_ids"], skip_special_tokens=False) + for decoded, assistant in zip(decoded_batch, assistant_texts, strict=True): + assert assistant in decoded + + # Verify ULD cross-tokenizer distillation with teacher inputs + student_tokenizer = smolvlm_processor.tokenizer + teacher_tokenizer = qwen3_vl_processor.tokenizer + + teacher_input_ids, teacher_labels, completion_texts = _teacher_inputs_from_collator( + student_tokenizer, teacher_tokenizer, batch + ) + for completion, assistant in zip(completion_texts, assistant_texts, strict=True): + assert assistant.strip() in completion + assert completion.strip() + + config = build_config( + uld_use_hybrid_loss=True, + uld_hybrid_matched_weight=0.6, + uld_hybrid_unmatched_weight=0.4, + ) + loss_fn = ULDLoss(config, student_tokenizer=student_tokenizer, teacher_tokenizer=teacher_tokenizer) + + _assert_alignment_covers_completion(loss_fn, batch, teacher_input_ids, teacher_labels) + + torch.manual_seed(42) + student_vocab = len(student_tokenizer) + teacher_vocab = len(teacher_tokenizer) + batch_size, seq_len = batch["input_ids"].shape + student_logits = torch.randn(batch_size, seq_len, student_vocab) + teacher_logits = torch.randn(batch_size, teacher_input_ids.shape[1], teacher_vocab) + + loss = loss_fn( + student_logits=student_logits, + teacher_logits=teacher_logits, + student_labels=batch["labels"], + teacher_labels=teacher_labels, + student_input_ids=batch["input_ids"], + teacher_input_ids=teacher_input_ids, + ) + + assert torch.isfinite(loss) + + +@pytest.mark.slow +def test_vlm_chatml_collator_preserves_completion_qwen3vl(smolvlm_processor, qwen3_vl_processor, vlm_examples): + collator = DataCollatorForVisionLanguageChatML(processor=qwen3_vl_processor, max_length=2048) + batch = collator(vlm_examples) + + # Verify basic batch structure + assert "input_ids" in batch + assert "labels" in batch + assert "prompts" in batch + assert "pixel_values" in batch + + # Verify completions are preserved in decoded output + assistant_texts = _get_assistant_texts(vlm_examples) + decoded_batch = qwen3_vl_processor.tokenizer.batch_decode(batch["input_ids"], skip_special_tokens=False) + for decoded, assistant in zip(decoded_batch, assistant_texts, strict=True): + assert assistant in decoded + + # Verify ULD cross-tokenizer distillation with teacher inputs + student_tokenizer = qwen3_vl_processor.tokenizer + teacher_tokenizer = smolvlm_processor.tokenizer + + teacher_input_ids, teacher_labels, completion_texts = _teacher_inputs_from_collator( + student_tokenizer, teacher_tokenizer, batch + ) + for completion, assistant in zip(completion_texts, assistant_texts, strict=True): + assert assistant.strip() in completion + assert completion.strip() + + config = build_config( + uld_use_hybrid_loss=True, + uld_hybrid_matched_weight=0.6, + uld_hybrid_unmatched_weight=0.4, + ) + loss_fn = ULDLoss(config, student_tokenizer=student_tokenizer, teacher_tokenizer=teacher_tokenizer) + + _assert_alignment_covers_completion(loss_fn, batch, teacher_input_ids, teacher_labels) + + torch.manual_seed(43) + student_vocab = len(student_tokenizer) + teacher_vocab = len(teacher_tokenizer) + batch_size, seq_len = batch["input_ids"].shape + student_logits = torch.randn(batch_size, seq_len, student_vocab) + teacher_logits = torch.randn(batch_size, teacher_input_ids.shape[1], teacher_vocab) + + loss = loss_fn( + student_logits=student_logits, + teacher_logits=teacher_logits, + student_labels=batch["labels"], + teacher_labels=teacher_labels, + student_input_ids=batch["input_ids"], + teacher_input_ids=teacher_input_ids, + ) + + assert torch.isfinite(loss) + + +def test_vlm_collator_label_masking(smolvlm_processor, vlm_examples): + """Verify that the VLM collator masks prompt and padding tokens in labels and leaves completion tokens unmasked.""" + collator = DataCollatorForVisionLanguageChatML(processor=smolvlm_processor, max_length=2048) + batch = collator(vlm_examples) + + input_ids = batch["input_ids"] + labels = batch["labels"] + attention_mask = batch["attention_mask"] + + for i in range(input_ids.shape[0]): + # Padding tokens (attention_mask == 0) must be masked in labels + padding_positions = attention_mask[i] == 0 + assert (labels[i][padding_positions] == -100).all(), "Padding tokens should be masked with -100" + + # There must be at least one non-masked label (completion token) + completion_positions = labels[i] != -100 + assert completion_positions.any(), "Each example must have at least one completion token in labels" + + # Completion labels must match the corresponding input_ids + assert (labels[i][completion_positions] == input_ids[i][completion_positions]).all(), ( + "Unmasked labels must match input_ids" + ) + + # Prompt tokens (attended but masked in labels) must exist — the prompt is never empty + prompt_positions = (attention_mask[i] == 1) & (labels[i] == -100) + assert prompt_positions.any(), "Each example must have masked prompt tokens" + + +def test_gold_trainer_init_rejects_non_vlm_teacher(monkeypatch): + """GOLDTrainer should raise ValueError when the student is a VLM but the teacher is not.""" + + class DummyStudentModel: + def __init__(self): + self.config = SimpleNamespace(_name_or_path="student", vocab_size=17) + self.generation_config = SimpleNamespace(eos_token_id=2) + self.name_or_path = "student" + + class DummyTeacherModel: + def __init__(self): + # No vision_config — looks like a text-only model + self.config = SimpleNamespace() + self.resized_to = None + + def resize_token_embeddings(self, vocab_size): + self.resized_to = vocab_size + + def fake_sft_init( + self, + model, + args=None, + data_collator=None, + train_dataset=None, + eval_dataset=None, + processing_class=None, + compute_metrics=None, + callbacks=None, + optimizers=None, + preprocess_logits_for_metrics=None, + peft_config=None, + ): + del data_collator, train_dataset, eval_dataset, compute_metrics, callbacks, optimizers + del preprocess_logits_for_metrics, peft_config + self.model = model + self.args = args + self.processing_class = processing_class + self.accelerator = SimpleNamespace( + device=torch.device("cpu"), + num_processes=1, + prepare_model=lambda module, evaluation_mode=True: module, + ) + self.is_deepspeed_enabled = False + self.is_fsdp_enabled = False + + monkeypatch.setattr(gold_trainer_module.SFTTrainer, "__init__", fake_sft_init) + + processor = AutoProcessor.from_pretrained("HuggingFaceTB/SmolVLM-256M-Instruct") + + vision_dataset = Dataset.from_dict({"messages": [["dummy"]], "image": ["fake_image"]}) + + args = SimpleNamespace( + model_init_kwargs=None, + max_length=128, + use_liger_kernel=False, + teacher_model_init_kwargs=None, + use_uld_loss=False, + teacher_tokenizer_name_or_path=None, + teacher_model_revision=None, + disable_dropout=False, + lmbda=1.0, + beta=0.5, + temperature=1.0, + top_p=1.0, + seq_kd=False, + num_generations=1, + use_transformers_paged=False, + max_completion_length=16, + top_k=0, + log_completions=False, + log_completions_steps=100, + wandb_log_unique_prompts=True, + num_completions_to_print=None, + per_device_train_batch_size=1, + gradient_accumulation_steps=1, + use_vllm=False, + ) + + with pytest.raises(ValueError, match="VLM distillation requires both student and teacher"): + GOLDTrainer( + model=DummyStudentModel(), + teacher_model=DummyTeacherModel(), + args=args, + train_dataset=vision_dataset, + processing_class=processor, + ) diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index caaeb9cdc9e..b8072130e3b 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -29,7 +29,7 @@ from accelerate.utils import DistributedType, broadcast_object_list, gather_object from datasets import Dataset, IterableDataset from torch.utils.data import DataLoader -from transformers import AutoTokenizer, TrainerCallback +from transformers import AutoProcessor, AutoTokenizer, TrainerCallback from transformers.data.data_collator import DataCollator from transformers.feature_extraction_utils import FeatureExtractionMixin from transformers.generation.configuration_utils import GenerationConfig @@ -58,10 +58,11 @@ RepeatSampler, create_model_from_path, disable_dropout_in_model, + get_config_model_id, pad, split_tensor_dict, ) -from ..utils import DataCollatorForChatML, empty_cache, truncate_dataset +from ..utils import DataCollatorForChatML, DataCollatorForVisionLanguageChatML, empty_cache, truncate_dataset from .gold_config import GOLDConfig @@ -778,10 +779,51 @@ def __init__( ): self.model_name_or_path = model if isinstance(model, str) else model.config._name_or_path self.model_revision = (args.model_init_kwargs or {}).get("revision") + if train_dataset is None: + raise ValueError("`train_dataset` is required") + dataset_sample = next(iter(train_dataset)) + if processing_class is None: + processing_class = AutoProcessor.from_pretrained(get_config_model_id(model.config)) + # simplified logic from SFTTrainer + if isinstance(processing_class, ProcessorMixin): + self._is_vlm = True + else: + self._is_vlm = False + + # VLM distillation: only VLM-to-VLM is supported. Both student and teacher must be + # VLMs so that both receive images and multimodal inputs. + if self._is_vlm and isinstance(teacher_model, str): + # Teacher not yet instantiated + teacher_proc = AutoProcessor.from_pretrained(teacher_model) + if not isinstance(teacher_proc, ProcessorMixin): + raise ValueError( + "VLM distillation requires both student and teacher to be vision-language models. " + "The student has a `ProcessorMixin` but the teacher does not." + ) + elif self._is_vlm and not isinstance(teacher_model, str): + # Teacher already instantiated — check if it looks like a VLM by checking for a vision config + if not hasattr(teacher_model, "config") or not hasattr(teacher_model.config, "vision_config"): + raise ValueError( + "VLM distillation requires both student and teacher to be vision-language models. " + "The student has a `ProcessorMixin` but the teacher model does not appear to be a VLM " + "(missing `vision_config`)." + ) + self._is_vision_dataset = "image" in dataset_sample or "images" in dataset_sample + if self._is_vision_dataset and not self._is_vlm: + raise ValueError( + "The dataset appears to be vision-related (contains 'image' or 'images' keys), but the provided " + "model does not seem to be a vision-language model. Please check your model and dataset." + ) - # Respect a user-provided data_collator; otherwise, provide a ChatML collator that + # Respect a user-provided data_collator; otherwise, pick the right collator based on modality if data_collator is None: - data_collator = DataCollatorForChatML(tokenizer=processing_class, max_length=args.max_length) + if self._is_vision_dataset: + data_collator = DataCollatorForVisionLanguageChatML( + processor=processing_class, + max_length=args.max_length, + ) + else: + data_collator = DataCollatorForChatML(tokenizer=processing_class, max_length=args.max_length) # Liger fused GKD loss (JSD) self.use_liger_gkd_loss = False @@ -974,6 +1016,15 @@ def _set_signature_columns_if_needed(self): "tools", "original_prompt_text", "original_completion_text", + "images", + "image", + "pixel_values", + "image_grid_thw", + "image_position_ids", + "pixel_attention_mask", + "image_sizes", + "token_type_ids", + "mm_token_type_ids", ] if self._signature_columns is None: self._signature_columns = required_columns @@ -1190,6 +1241,9 @@ def _generate_on_policy_for_slices( self.vllm_generation.sync_weights() self._last_vllm_sync_step = self.state.global_step + # TODO: pass raw images from the dataset to vLLM for VLM on-policy generation. + # Currently, the collated batch only contains processed pixel_values, not raw images. + # vLLM generation with VLMs requires raw images which are lost after collation. _, completion_ids, _, _ = self.vllm_generation.generate( prompts=prompt_ids_list, images=None, @@ -1367,6 +1421,11 @@ def _prepare_dataset( dataset_name: str, ) -> Dataset | IterableDataset: """Preserve original text fields for ULD when needed.""" + # For VLM datasets, skip dataset preparation entirely — the VLM collator handles tokenization + # and image processing on the fly, similar to how SFTTrainer skips prep for vision datasets. + if self._is_vision_dataset: + return dataset + column_names = list(next(iter(dataset)).keys()) is_processed = "input_ids" in column_names @@ -1689,7 +1748,20 @@ def generalized_jsd_loss( else: return jsd + _MULTIMODAL_KEYS = ( + "pixel_values", + "image_grid_thw", + "image_position_ids", + "pixel_attention_mask", + "image_sizes", + "token_type_ids", + "mm_token_type_ids", + ) + def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): + # Extract multimodal fields for VLM forward passes + forward_kwargs = {k: inputs[k] for k in self._MULTIMODAL_KEYS if k in inputs} + if self.use_uld_loss and self.teacher_tokenizer is not None: if "original_prompt_text" in inputs and "original_completion_text" in inputs: prompt_texts = inputs["original_prompt_text"] @@ -1727,6 +1799,7 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"], use_cache=False, + **forward_kwargs, ) self.teacher_model.eval() @@ -1734,6 +1807,7 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N outputs_teacher = self.teacher_model( input_ids=teacher_input_ids, attention_mask=teacher_attention_mask, + **forward_kwargs, ) # These are not used for ULD loss but are needed if JSD loss were to be used in this branch @@ -1756,6 +1830,7 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"], use_cache=False, + **forward_kwargs, ) self.teacher_model.eval() @@ -1771,6 +1846,7 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"], use_cache=False, + **forward_kwargs, ) student_hidden = student_outputs.last_hidden_state[:, :-1] @@ -1805,6 +1881,7 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N outputs_student = model( input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"], + **forward_kwargs, ) self.teacher_model.eval() @@ -1812,6 +1889,7 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N outputs_teacher = self.teacher_model( input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"], + **forward_kwargs, ) prompt_lengths = inputs["prompts"].shape[1] @@ -1898,11 +1976,13 @@ def generate_on_policy_outputs(self, model, inputs, generation_config, pad_token completion_ids = [output.generated_tokens for output in generated_outputs.values()] generated_tokens = torch.stack([torch.tensor(ids, device=model.device) for ids in completion_ids]) else: + generate_kwargs = {k: inputs[k] for k in self._MULTIMODAL_KEYS if k in inputs} generated_outputs = model.generate( input_ids=inputs["prompts"], attention_mask=inputs.get("prompt_attention_mask", None), generation_config=generation_config, return_dict_in_generate=True, + **generate_kwargs, ) # Get the generated token IDs generated_tokens = generated_outputs.sequences diff --git a/trl/experimental/utils.py b/trl/experimental/utils.py index c338283412f..8115342518d 100644 --- a/trl/experimental/utils.py +++ b/trl/experimental/utils.py @@ -29,7 +29,9 @@ from torch import nn from torch.nn.utils.rnn import pad_sequence from transformers import PreTrainedModel, PreTrainedTokenizerBase, TrainingArguments +from transformers.data.data_collator import DataCollatorMixin from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled +from transformers.processing_utils import ProcessorMixin from transformers.utils import ( is_peft_available, is_torch_mlu_available, @@ -37,8 +39,14 @@ is_torch_xpu_available, ) -from ..data_utils import DatasetType, _get_dataset_format -from ..trainer.utils import pad +from ..data_utils import ( + DatasetType, + _get_dataset_format, + apply_chat_template, + is_conversational, + prepare_multimodal_messages, +) +from ..trainer.utils import flush_left, pad if is_peft_available(): @@ -261,6 +269,155 @@ def __call__(self, examples: list[dict[str, Any]]) -> dict[str, torch.Tensor]: } +@dataclass +class DataCollatorForVisionLanguageChatML(DataCollatorMixin): + """ + Data collator for GOLD VLM training. + + Combines image processing from [`~trainer.sft_trainer.DataCollatorForVisionLanguageModeling`] with the + prompt-separation logic that GOLD needs for on-policy generation. Each input example should be a dictionary + containing at least: + - An `"images"` key holding a list of images, or an `"image"` key holding a single image. + - Keys `"prompt"` and `"completion"` for the prompt and completion (conversational or plain text). + + The collator outputs a dictionary including: + - `"input_ids"`: Tensor of token IDs (prompt + completion, concatenated). + - `"attention_mask"`: Tensor indicating attention mask. + - `"labels"`: Tensor for training labels (prompt tokens masked with -100). + - `"prompts"`: Tensor of prompt-only token IDs (left-padded), used for on-policy generation. + - `"prompt_attention_mask"`: Attention mask for prompts. + - `"original_prompt_text"`: List of raw prompt text strings, used for ULD cross-tokenizer distillation. + - `"original_completion_text"`: List of raw completion text strings, used for ULD cross-tokenizer distillation. + - `"pixel_values"`: Tensor representing image pixel values. + + Additional keys may be present depending on the processor, such as `"image_grid_thw"` or `"image_position_ids"`. + + Args: + processor ([`~transformers.ProcessorMixin`]): + The processor used to tokenize text and process images. + max_length (`int` or `None`, *optional*): + Maximum sequence length for input tokens. If `None`, no truncation is applied. + return_tensors (`str`, *optional*, defaults to `"pt"`): + The tensor type to return. + """ + + processor: ProcessorMixin + max_length: int | None = None + return_tensors: str = "pt" + + def torch_call(self, examples: list[dict[str, Any]]) -> dict[str, Any]: + if "prompt" not in examples[0] or "completion" not in examples[0]: + raise KeyError( + "DataCollatorForVisionLanguageChatML requires 'prompt' and 'completion' keys in examples. " + f"Got keys: {list(examples[0].keys())}." + ) + + # Normalize single image to list + if "image" in examples[0]: + for example in examples: + example["images"] = [example.pop("image")] + images = [example.get("images", []) for example in examples] + if all(img_list == [] for img_list in images): + images = None + + # Apply chat template for conversational data + if is_conversational(examples[0]): + for example in examples: + example["prompt"] = prepare_multimodal_messages(example["prompt"], images=example["images"]) + example["completion"] = prepare_multimodal_messages(example["completion"]) + examples = [apply_chat_template(example, self.processor) for example in examples] + + prompts = [example["prompt"] for example in examples] + completions = [example["completion"] for example in examples] + + # Process prompts (with images) and completions (text only) separately + processed_prompts = self.processor( + images=images, + text=prompts, + padding=True, + padding_side="left", + return_tensors=self.return_tensors, + add_special_tokens=False, + ) + processed_completions = self.processor( + text=completions, + padding=True, + padding_side="right", + return_tensors=self.return_tensors, + add_special_tokens=False, + ) + + # Concatenate prompts and completions + prompt_ids, prompt_mask = processed_prompts["input_ids"], processed_prompts["attention_mask"] + completion_ids, completion_mask = processed_completions["input_ids"], processed_completions["attention_mask"] + input_ids = torch.cat((prompt_ids, completion_ids), dim=1) + attention_mask = torch.cat((prompt_mask, completion_mask), dim=1) + completion_mask = torch.cat((torch.zeros_like(prompt_mask), completion_mask), dim=1) + if "token_type_ids" in processed_prompts: + prompt_token_type_ids = processed_prompts["token_type_ids"] + completion_token_type_ids = processed_completions["token_type_ids"] + token_type_ids = torch.cat((prompt_token_type_ids, completion_token_type_ids), dim=1) + if "mm_token_type_ids" in processed_prompts: + prompt_mm_token_type_ids = processed_prompts["mm_token_type_ids"] + completion_mm_token_type_ids = processed_completions.get( + "mm_token_type_ids", torch.zeros_like(completion_ids) + ) + mm_token_type_ids = torch.cat((prompt_mm_token_type_ids, completion_mm_token_type_ids), dim=1) + + # Flush left to reduce padding + if "token_type_ids" in processed_prompts and "mm_token_type_ids" in processed_prompts: + attention_mask, input_ids, completion_mask, token_type_ids, mm_token_type_ids = flush_left( + attention_mask, input_ids, completion_mask, token_type_ids, mm_token_type_ids + ) + elif "token_type_ids" in processed_prompts: + attention_mask, input_ids, completion_mask, token_type_ids = flush_left( + attention_mask, input_ids, completion_mask, token_type_ids + ) + elif "mm_token_type_ids" in processed_prompts: + attention_mask, input_ids, completion_mask, mm_token_type_ids = flush_left( + attention_mask, input_ids, completion_mask, mm_token_type_ids + ) + else: + attention_mask, input_ids, completion_mask = flush_left(attention_mask, input_ids, completion_mask) + + # Truncate if necessary + if self.max_length is not None: + input_ids = input_ids[:, : self.max_length] + attention_mask = attention_mask[:, : self.max_length] + completion_mask = completion_mask[:, : self.max_length] + if "token_type_ids" in processed_prompts: + token_type_ids = token_type_ids[:, : self.max_length] + if "mm_token_type_ids" in processed_prompts: + mm_token_type_ids = mm_token_type_ids[:, : self.max_length] + + # Create labels: mask padding and prompt tokens + labels = input_ids.clone() + labels[attention_mask == 0] = -100 + labels[completion_mask == 0] = -100 + + # Build output with vision keys from processed_prompts (pixel_values, image_grid_thw, etc.) + output = processed_prompts + output["input_ids"] = input_ids + output["attention_mask"] = attention_mask + output["labels"] = labels + if "token_type_ids" in processed_prompts: + output["token_type_ids"] = token_type_ids + if ( + "mm_token_type_ids" in processed_prompts + ): # special case for ERNIE-VL from class DataCollatorForVisionLanguageModeling(DataCollatorMixin): + output["mm_token_type_ids"] = mm_token_type_ids + + # GOLD-specific: separate prompt tensors for on-policy generation + output["prompts"] = prompt_ids + output["prompt_attention_mask"] = prompt_mask + + # GOLD-specific: raw text for ULD cross-tokenizer distillation + output["original_prompt_text"] = prompts + output["original_completion_text"] = completions + + return output + + def truncate_right( input_ids: torch.Tensor, stop_token_id: int, pad_token_id: int ) -> tuple[torch.Tensor, torch.Tensor]: From a92473bd09f2076f489042c31d120db6403e99f7 Mon Sep 17 00:00:00 2001 From: Strongich Date: Sun, 5 Apr 2026 21:24:07 +0200 Subject: [PATCH 02/39] add vlm support with use_vllm=True --- tests/experimental/test_gold_trainer.py | 126 +++++++++++++++ trl/experimental/gold/gold_trainer.py | 199 +++++++++++++++++++++--- 2 files changed, 307 insertions(+), 18 deletions(-) diff --git a/tests/experimental/test_gold_trainer.py b/tests/experimental/test_gold_trainer.py index 3257ad2f99f..ace0c4ac1e7 100644 --- a/tests/experimental/test_gold_trainer.py +++ b/tests/experimental/test_gold_trainer.py @@ -22,6 +22,7 @@ from trl.experimental.gold import gold_trainer as gold_trainer_module from trl.experimental.gold.gold_trainer import GOLDTrainer, ULDLoss, build_teacher_inputs_from_texts from trl.experimental.utils import DataCollatorForChatML, DataCollatorForVisionLanguageChatML +from trl.trainer.utils import identity @pytest.fixture(scope="module") @@ -331,6 +332,7 @@ def batch_decode(self, sequences, skip_special_tokens=False, clean_up_tokenizati trainer = GOLDTrainer.__new__(GOLDTrainer) trainer.accelerator = SimpleNamespace(device=torch.device("cpu")) trainer.processing_class = RecordingTokenizer() + trainer.pad_token_id = RecordingTokenizer.pad_token_id # __new__ bypasses __init__, set manually trainer.args = SimpleNamespace(max_length=None) trainer._buffered_inputs = [None] trainer._buffered_text_logs = [None] @@ -478,6 +480,7 @@ def batch_decode(self, sequences, skip_special_tokens=False, clean_up_tokenizati trainer = GOLDTrainer.__new__(GOLDTrainer) trainer.accelerator = SimpleNamespace(device=torch.device("cpu"), is_main_process=True) trainer.processing_class = RecordingTokenizer() + trainer.pad_token_id = RecordingTokenizer.pad_token_id # __new__ bypasses __init__, set manually trainer.args = SimpleNamespace(max_length=None, report_to=[]) trainer.use_vllm = True trainer.vllm_generation = RecordingVLLMGeneration() @@ -527,6 +530,9 @@ def resize_token_embeddings(self, vocab_size): self.resized_to = vocab_size class DummyProcessingClass: + # GOLDTrainer.__init__ extracts tokenizer pad token (like GRPOTrainer), + # so the dummy must provide both pad_token and pad_token_id. + pad_token = "" pad_token_id = 0 def fake_sft_init( @@ -1328,3 +1334,123 @@ def fake_sft_init( train_dataset=vision_dataset, processing_class=processor, ) + + +def test_gold_trainer_vlm_vllm_init_uses_identity_collator(monkeypatch): + """When a VLM processor is used with lmbda > 0 and use_vllm=True, GOLDTrainer should use the identity collator + and store a _vlm_collator for on-the-fly collation. vLLM should be initialized with max_model_length from args.""" + captured = {} + + class DummyStudentModel: + def __init__(self): + self.config = SimpleNamespace(_name_or_path="student", vocab_size=17, vision_config=True) + self.generation_config = SimpleNamespace(eos_token_id=2) + self.name_or_path = "student" + + class DummyTeacherModel: + def __init__(self): + self.config = SimpleNamespace(vision_config=True) + self.resized_to = None + + def resize_token_embeddings(self, vocab_size): + self.resized_to = vocab_size + + def fake_sft_init( + self, + model, + args=None, + data_collator=None, + train_dataset=None, + eval_dataset=None, + processing_class=None, + compute_metrics=None, + callbacks=None, + optimizers=None, + preprocess_logits_for_metrics=None, + peft_config=None, + ): + self.data_collator = data_collator + del train_dataset, eval_dataset, compute_metrics, callbacks, optimizers + del preprocess_logits_for_metrics, peft_config + self.model = model + self.args = args + self.processing_class = processing_class + self.accelerator = SimpleNamespace( + device=torch.device("cpu"), + num_processes=1, + prepare_model=lambda module, evaluation_mode=True: module, + ) + self.is_deepspeed_enabled = False + self.is_fsdp_enabled = False + + class CapturingVLLMGeneration: + def __init__(self, **kwargs): + captured.update(kwargs) + + monkeypatch.setattr(gold_trainer_module.SFTTrainer, "__init__", fake_sft_init) + monkeypatch.setattr(gold_trainer_module, "is_vllm_available", lambda: True) + monkeypatch.setattr(gold_trainer_module, "VLLMGeneration", CapturingVLLMGeneration) + + processor = AutoProcessor.from_pretrained("HuggingFaceTB/SmolVLM-256M-Instruct") + if processor.tokenizer.pad_token is None: + processor.tokenizer.pad_token = processor.tokenizer.eos_token + + vision_dataset = Dataset.from_dict({"messages": [["dummy"]], "image": ["fake_image"]}) + + args = SimpleNamespace( + model_init_kwargs=None, + max_length=128, + use_liger_kernel=False, + teacher_model_init_kwargs=None, + use_uld_loss=False, + teacher_tokenizer_name_or_path=None, + teacher_model_revision=None, + disable_dropout=False, + lmbda=1.0, + beta=0.5, + temperature=1.0, + top_p=1.0, + seq_kd=False, + num_generations=1, + use_transformers_paged=False, + max_completion_length=16, + top_k=0, + log_completions=False, + log_completions_steps=100, + wandb_log_unique_prompts=True, + num_completions_to_print=None, + per_device_train_batch_size=1, + gradient_accumulation_steps=1, + use_vllm=True, + vllm_mode="colocate", + vllm_structured_outputs_regex=None, + vllm_server_base_url=None, + vllm_server_host="0.0.0.0", + vllm_server_port=8001, + vllm_group_port=51216, + vllm_server_timeout=240.0, + vllm_tensor_parallel_size=1, + vllm_gpu_memory_utilization=0.2, + vllm_max_model_length=None, + vllm_enable_sleep_mode=False, + vllm_model_impl="vllm", + vllm_sync_frequency=1, + ) + + teacher_model = DummyTeacherModel() + trainer = GOLDTrainer( + model=DummyStudentModel(), + teacher_model=teacher_model, + args=args, + train_dataset=vision_dataset, + processing_class=processor, + ) + + # Same assertions as text-only vLLM test + assert teacher_model.resized_to == 17 + assert captured["max_model_length"] == 128 + + # VLM-specific: identity collator + _vlm_collator for on-the-fly use + assert trainer.data_collator is identity + assert trainer._vlm_collator is not None + assert isinstance(trainer._vlm_collator, DataCollatorForVisionLanguageChatML) diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index b8072130e3b..933189a0970 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -47,7 +47,12 @@ is_rich_available, ) -from ...data_utils import is_conversational, maybe_convert_to_chatml, pack_dataset +from ...data_utils import ( + is_conversational, + maybe_convert_to_chatml, + pack_dataset, + prepare_multimodal_messages, +) from ...extras.profiling import profiling_decorator from ...generation.vllm_generation import VLLMGeneration from ...import_utils import is_vllm_available @@ -59,6 +64,7 @@ create_model_from_path, disable_dropout_in_model, get_config_model_id, + identity, pad, split_tensor_dict, ) @@ -785,11 +791,19 @@ def __init__( if processing_class is None: processing_class = AutoProcessor.from_pretrained(get_config_model_id(model.config)) # simplified logic from SFTTrainer + # Handle pad token for processors or tokenizers if isinstance(processing_class, ProcessorMixin): + tokenizer = processing_class.tokenizer self._is_vlm = True else: + tokenizer = processing_class self._is_vlm = False + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + + self.pad_token_id = tokenizer.pad_token_id + # VLM distillation: only VLM-to-VLM is supported. Both student and teacher must be # VLMs so that both receive images and multimodal inputs. if self._is_vlm and isinstance(teacher_model, str): @@ -815,9 +829,19 @@ def __init__( "model does not seem to be a vision-language model. Please check your model and dataset." ) - # Respect a user-provided data_collator; otherwise, pick the right collator based on modality + # Respect a user-provided data_collator; otherwise, pick the right collator based on modality. + # For VLMs with lmbda > 0, use identity collator to preserver raw PIL images in the dataloader. + # Because vLLM requires raw images, not processed pixel_values tensors. + # A separate _vlm_collator is stored for on-the-fly collation inside _fill_buffer. + self._vlm_collator = None if data_collator is None: - if self._is_vision_dataset: + if self._is_vision_dataset and args.lmbda > 0: + self._vlm_collator = DataCollatorForVisionLanguageChatML( + processor=processing_class, + max_length=args.max_length, + ) + data_collator = identity + elif self._is_vision_dataset: data_collator = DataCollatorForVisionLanguageChatML( processor=processing_class, max_length=args.max_length, @@ -942,7 +966,7 @@ def __init__( "top_p": args.top_p, "do_sample": True, "top_k": args.top_k, - "pad_token_id": self.processing_class.pad_token_id, + "pad_token_id": self.pad_token_id, } self.generation_config = GenerationConfig(**generation_kwargs) # Keep training-specific generation kwargs to overwrite model's original generation config @@ -1109,8 +1133,8 @@ def _decode_completion_texts_from_labels(self, slice_inputs: dict[str, torch.Ten decoded_completion_tokens: list[list[int]] = [] for row in labels_cpu: token_ids = row[row != -100].tolist() - if self.processing_class.pad_token_id is not None: - token_ids = [tok for tok in token_ids if tok != self.processing_class.pad_token_id] + if self.pad_token_id is not None: + token_ids = [tok for tok in token_ids if tok != self.pad_token_id] decoded_completion_tokens.append(token_ids) return self.processing_class.batch_decode( @@ -1173,8 +1197,16 @@ def _build_sequence_batch( return new_attention_mask, new_labels @profiling_decorator - def _fill_buffer(self, generation_batch: dict[str, torch.Tensor | Any], buffer_steps: int): - slices = split_tensor_dict(generation_batch, buffer_steps) + def _fill_buffer(self, generation_batch: dict[str, torch.Tensor | Any] | list[dict], buffer_steps: int): + if self._vlm_collator is not None: + # Identity collator path: generation_batch is list[dict] with raw PIL images. + # Split into chunks via list slicing, then collate on-the-fly per slice. + chunk_size = len(generation_batch) // buffer_steps + raw_slices = [generation_batch[i * chunk_size : (i + 1) * chunk_size] for i in range(buffer_steps)] + slices = None # not used in this path + else: + raw_slices = None # not used in this path + slices = split_tensor_dict(generation_batch, buffer_steps) if self.accelerator.is_main_process: on_policy_flags = [random.random() <= self.lmbda for _ in range(buffer_steps)] @@ -1190,7 +1222,15 @@ def _fill_buffer(self, generation_batch: dict[str, torch.Tensor | Any], buffer_s for i, flag in enumerate(on_policy_flags): if not flag: - slice_inputs = slices[i] + if self._vlm_collator is not None: + # Collate raw examples on-the-fly for off-policy slices + slice_inputs = self._vlm_collator(raw_slices[i]) + slice_inputs = { + k: v.to(self.accelerator.device) if isinstance(v, torch.Tensor) else v + for k, v in slice_inputs.items() + } + else: + slice_inputs = slices[i] if self.use_uld_loss and self.teacher_tokenizer is not None: slice_inputs = self._ensure_original_text_fields(slice_inputs) @@ -1204,7 +1244,10 @@ def _fill_buffer(self, generation_batch: dict[str, torch.Tensor | Any], buffer_s self._buffered_inputs[i] = slice_inputs if on_policy_indices: - self._generate_on_policy_for_slices(slices, on_policy_indices) + if self._vlm_collator is not None: + self._generate_on_policy_vlm_raw(raw_slices, on_policy_indices) + else: + self._generate_on_policy_for_slices(slices, on_policy_indices) @profiling_decorator def _generate_on_policy_for_slices( @@ -1241,9 +1284,8 @@ def _generate_on_policy_for_slices( self.vllm_generation.sync_weights() self._last_vllm_sync_step = self.state.global_step - # TODO: pass raw images from the dataset to vLLM for VLM on-policy generation. - # Currently, the collated batch only contains processed pixel_values, not raw images. - # vLLM generation with VLMs requires raw images which are lost after collation. + # Text-only vLLM generation. VLM on-policy generation with raw images + # is handled by _generate_on_policy_vlm_raw (routed from _fill_buffer). _, completion_ids, _, _ = self.vllm_generation.generate( prompts=prompt_ids_list, images=None, @@ -1274,7 +1316,7 @@ def _generate_non_vllm_for_slices(self, slices: list[dict[str, torch.Tensor | An unwrapped_model, slice_inputs, self.generation_config, - self.processing_class.pad_token_id, + self.pad_token_id, ) new_input_ids, new_attention_mask, new_labels, prompt_texts, completion_texts = result @@ -1288,6 +1330,127 @@ def _generate_non_vllm_for_slices(self, slices: list[dict[str, torch.Tensor | An self._buffered_inputs[slice_idx] = updated_slice self._buffered_text_logs[slice_idx] = (prompt_texts, completion_texts) + def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_indices: list[int]): + """On-policy generation from raw VLM examples, preserving PIL images for vLLM.""" + device = self.accelerator.device + + for slice_idx in on_policy_indices: + raw_examples = raw_slices[slice_idx] + + # Extract raw PIL images from examples (like GRPOTrainer) + if "images" in raw_examples[0]: + images = [example.get("images") for example in raw_examples] + elif "image" in raw_examples[0]: + images = [ + [example.get("image")] if example.get("image") is not None else None for example in raw_examples + ] + else: + images = None + if images is not None and all(img_list is None or img_list == [] for img_list in images): + images = None + + # Extract prompts and prepare multimodal messages + prompts = [example["prompt"] for example in raw_examples] + if images is not None: + prompts = [ + prepare_multimodal_messages(prompt, images=img_list) + for prompt, img_list in zip(prompts, images, strict=True) + ] + + # Normalize string content to content blocks for VLM processors that don't handle plain strings + # copied from GRPOTrainer + prompts = [ + [ + {**msg, "content": [{"type": "text", "text": msg["content"]}]} + if isinstance(msg.get("content"), str) + else msg + for msg in prompt + ] + for prompt in prompts + ] + + # Tokenize prompts to get prompt token IDs + # TODO: add self.tools support + tokenized = self.processing_class.apply_chat_template( + conversation=prompts, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + padding=True, + ) + prompt_ids_list = [ + [tok for tok, m in zip(ids, mask, strict=True) if m] + for ids, mask in zip(tokenized["input_ids"], tokenized["attention_mask"], strict=True) + ] + + prompts_text = self.processing_class.batch_decode(prompt_ids_list, skip_special_tokens=True) + prompts_text_with_special = self.processing_class.batch_decode(prompt_ids_list, skip_special_tokens=False) + + if not self.use_vllm: + # Non-vLLM path: collate raw examples to get pixel_values, then generate + collated = self._vlm_collator(raw_examples) + collated = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in collated.items()} + with unwrap_model_for_generation( + self.model, self.accelerator, generation_kwargs=self.generation_kwargs + ) as unwrapped_model: + result = self.generate_on_policy_outputs( + unwrapped_model, collated, self.generation_config, self.pad_token_id + ) + new_input_ids, new_attention_mask, new_labels, prompt_texts, completion_texts = result + + updated_slice = dict(collated) + updated_slice["input_ids"] = new_input_ids + updated_slice["attention_mask"] = new_attention_mask + updated_slice["labels"] = new_labels + updated_slice["original_prompt_text"] = prompt_texts + updated_slice["original_completion_text"] = completion_texts + + self._buffered_inputs[slice_idx] = updated_slice + self._buffered_text_logs[slice_idx] = (prompt_texts, completion_texts) + continue + + # vLLM path: pass raw PIL images to vLLM + if ( + self.state.global_step != self._last_vllm_sync_step + and self.state.global_step >= self._last_vllm_sync_step + self.vllm_sync_frequency + ): + self.vllm_generation.sync_weights() + self._last_vllm_sync_step = self.state.global_step + + _, completion_ids, _, _ = self.vllm_generation.generate( + prompts=prompt_ids_list, + images=images, + num_generations=self.num_generations, + ) + + # Decode completions and build synthetic examples for collation + max_completion_length = self.generation_config.max_new_tokens + completion_texts = [] + for comp_ids in completion_ids: + if len(comp_ids) > max_completion_length: + comp_ids = comp_ids[:max_completion_length] + completion_texts.append( + self.processing_class.decode( + comp_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False + ) + ) + + # Build synthetic examples: original prompt + generated completion + synthetic_examples = [] + for i, example in enumerate(raw_examples): + synthetic = dict(example) + synthetic["completion"] = [{"role": "assistant", "content": completion_texts[i]}] + synthetic_examples.append(synthetic) + + # Collate synthetic examples to get pixel_values + properly tokenized input_ids/labels + collated = self._vlm_collator(synthetic_examples) + collated = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in collated.items()} + collated["original_prompt_text"] = prompts_text_with_special + collated["original_completion_text"] = completion_texts + + self._buffered_inputs[slice_idx] = collated + self._buffered_text_logs[slice_idx] = (prompts_text, completion_texts) + def _process_completions_to_buffer( self, slices: list[dict[str, torch.Tensor | Any]], @@ -1303,7 +1466,7 @@ def _process_completions_to_buffer( Process vLLM completions and update buffered inputs for on-policy slices. """ device = self.accelerator.device - pad_token_id = self.processing_class.pad_token_id if self.processing_class.pad_token_id is not None else 0 + pad_token_id = self.pad_token_id if self.pad_token_id is not None else 0 slice_completions = {idx: [] for idx in on_policy_indices} slice_prompt_ids = {idx: [] for idx in on_policy_indices} @@ -1911,8 +2074,8 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N teacher_input_ids_for_loss = teacher_input_ids if "teacher_input_ids" in locals() else inputs["input_ids"] student_labels = inputs["labels"].clone() - if hasattr(self.processing_class, "pad_token_id") and self.processing_class.pad_token_id is not None: - student_labels[student_labels == self.processing_class.pad_token_id] = -100 + if self.pad_token_id is not None: + student_labels[student_labels == self.pad_token_id] = -100 if ( hasattr(self, "teacher_tokenizer") @@ -1991,7 +2154,7 @@ def generate_on_policy_outputs(self, model, inputs, generation_config, pad_token device = generated_tokens.device prompt_mask = inputs.get("prompt_attention_mask") - pad_token_id = pad_token_id if pad_token_id is not None else self.processing_class.pad_token_id + pad_token_id = pad_token_id if pad_token_id is not None else self.pad_token_id if self.use_transformers_paged: # generate_batch() returns completion-only tokens, so the entire tensor is completion. From 23784b095f5c6406db234f14d1df2f3c939ce8d2 Mon Sep 17 00:00:00 2001 From: Strongich Date: Mon, 6 Apr 2026 17:23:59 +0200 Subject: [PATCH 03/39] Add cross-architecture VLM distillation support to GOLDTrainer --- examples/scripts/gold_vlm.py | 181 +++++++++++++++ tests/experimental/test_gold_trainer.py | 293 +++++++++++++++++++++++- trl/experimental/gold/gold_trainer.py | 140 ++++++++--- 3 files changed, 584 insertions(+), 30 deletions(-) create mode 100644 examples/scripts/gold_vlm.py diff --git a/examples/scripts/gold_vlm.py b/examples/scripts/gold_vlm.py new file mode 100644 index 00000000000..f50ad67a691 --- /dev/null +++ b/examples/scripts/gold_vlm.py @@ -0,0 +1,181 @@ +# Copyright 2020-2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +GOLD VLM distillation on MMK12. + +# Example 1 — Same-family distillation (SmolVLM-500M → SmolVLM-256M) +# Uses JSD loss. Same architecture and tokenizer, so standard distillation works directly. +# vLLM enabled for faster on-policy generation. +accelerate launch examples/scripts/gold_vlm.py \ + --student_model_name HuggingFaceTB/SmolVLM-256M-Instruct \ + --teacher_model_name HuggingFaceTB/SmolVLM-500M-Instruct \ + --lmbda 0.5 \ + --use_vllm \ + --vllm_mode colocate + +# Example 2 — Cross-family distillation (Qwen2.5-VL-3B → SmolVLM-256M) +# Different architectures have incompatible tokenizers and image token formats, +# so ULD (Universal Logit Distillation) loss is required to align logits across vocabularies. +accelerate launch examples/scripts/gold_vlm.py \ + --student_model_name HuggingFaceTB/SmolVLM-256M-Instruct \ + --teacher_model_name Qwen/Qwen2.5-VL-3B-Instruct \ + --use_uld_loss \ + --lmbda 0.0 +""" + +import argparse + +import torch +from datasets import load_dataset +from peft import LoraConfig +from transformers import AutoModelForImageTextToText, AutoProcessor + +from trl.experimental.gold import GOLDConfig, GOLDTrainer + + +SYSTEM_PROMPT = ( + "You are a helpful AI Assistant that provides well-reasoned and detailed responses. " + "You first think about the reasoning process as an internal monologue and then provide the user with the answer. " + "Respond in the following format: \n...\n\n\n...\n" +) + + +def make_conversation(example): + """Convert MMK12 row into the chat format expected by TRL VLM trainers.""" + return { + "prompt": [ + { + "role": "system", + "content": [{"type": "text", "text": SYSTEM_PROMPT}], + }, + { + "role": "user", + "content": [ + {"type": "image"}, + {"type": "text", "text": example["question"]}, + ], + }, + ], + "completion": [ + { + "role": "assistant", + "content": [{"type": "text", "text": str(example["answer"])}], + }, + ], + "image": example["image"], + } + + +def filter_big_images(example): + image = example["image"] + return image.size[0] < 512 and image.size[1] < 512 + + +def convert_to_rgb(example): + image = example["image"] + if image.mode != "RGB": + image = image.convert("RGB") + example["image"] = image + return example + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--student_model_name", type=str, default="HuggingFaceTB/SmolVLM-256M-Instruct") + parser.add_argument("--teacher_model_name", type=str, default="HuggingFaceTB/SmolVLM-500M-Instruct") + parser.add_argument("--use_uld_loss", action="store_true") + parser.add_argument("--lmbda", type=float, default=0.5) + parser.add_argument("--use_vllm", action="store_true") + parser.add_argument("--vllm_mode", type=str, default="colocate") + cli_args = parser.parse_args() + + # ────────────────────────────────────────────── + # Models + # ────────────────────────────────────────────── + student_model = AutoModelForImageTextToText.from_pretrained(cli_args.student_model_name, dtype=torch.bfloat16) + teacher_model = AutoModelForImageTextToText.from_pretrained(cli_args.teacher_model_name, dtype=torch.bfloat16) + + # Freeze everything except the language model head + for name, param in student_model.named_parameters(): + if "language_model" not in name: + param.requires_grad = False + + processor = AutoProcessor.from_pretrained(cli_args.student_model_name, padding_side="left") + + # toy example to fit small GPUs + peft_config = LoraConfig( + r=4, + lora_alpha=8, + lora_dropout=0.05, + target_modules=["q_proj"], + ) + + # ────────────────────────────────────────────── + # Dataset + # ────────────────────────────────────────────── + dataset = load_dataset("FanqingM/MMK12", split="train[:5%]") + dataset = dataset.filter(filter_big_images) + dataset = dataset.map(convert_to_rgb) + dataset = dataset.map(make_conversation) + + # ────────────────────────────────────────────── + # Training config + # ────────────────────────────────────────────── + args = GOLDConfig( + output_dir="gold-vlm-distillation", + # GOLD-specific + lmbda=cli_args.lmbda, + beta=0.5, + temperature=0.9, + max_completion_length=256, + teacher_model_name_or_path=cli_args.teacher_model_name, + num_generations=1, + use_uld_loss=cli_args.use_uld_loss, + # vLLM + use_vllm=cli_args.use_vllm, + vllm_mode=cli_args.vllm_mode, + vllm_gpu_memory_utilization=0.5, + vllm_max_model_length=8192, + # VLM image tokens expand during processing, so the default max_length (1024) is often too small. + # Which will lead to shifted_student_logits become an empty Tensor. + max_length=2048, + # Training schedule + per_device_train_batch_size=2, + gradient_accumulation_steps=4, + max_steps=100, + learning_rate=2e-5, + warmup_steps=10, + # Precision + bf16=True, + # Logging + logging_steps=1, + log_completions=True, + report_to="none", + ) + + # ────────────────────────────────────────────── + # Trainer + # ────────────────────────────────────────────── + trainer = GOLDTrainer( + model=student_model, + teacher_model=teacher_model, + args=args, + train_dataset=dataset, + processing_class=processor, + peft_config=peft_config, + ) + + trainer.train() + trainer.save_model(args.output_dir) diff --git a/tests/experimental/test_gold_trainer.py b/tests/experimental/test_gold_trainer.py index ace0c4ac1e7..76df35cc03c 100644 --- a/tests/experimental/test_gold_trainer.py +++ b/tests/experimental/test_gold_trainer.py @@ -1343,13 +1343,15 @@ def test_gold_trainer_vlm_vllm_init_uses_identity_collator(monkeypatch): class DummyStudentModel: def __init__(self): - self.config = SimpleNamespace(_name_or_path="student", vocab_size=17, vision_config=True) + self.config = SimpleNamespace( + _name_or_path="student", vocab_size=17, vision_config=True, model_type="dummy_vlm" + ) self.generation_config = SimpleNamespace(eos_token_id=2) self.name_or_path = "student" class DummyTeacherModel: def __init__(self): - self.config = SimpleNamespace(vision_config=True) + self.config = SimpleNamespace(vision_config=True, model_type="dummy_vlm") self.resized_to = None def resize_token_embeddings(self, vocab_size): @@ -1454,3 +1456,290 @@ def __init__(self, **kwargs): assert trainer.data_collator is identity assert trainer._vlm_collator is not None assert isinstance(trainer._vlm_collator, DataCollatorForVisionLanguageChatML) + + +def _make_dummy_vlm_models(student_model_type, teacher_model_type): + """Helper to create dummy student/teacher VLM models with specified model_type.""" + + class DummyStudentModel: + def __init__(self): + self.config = SimpleNamespace( + _name_or_path="student", vocab_size=17, vision_config=True, model_type=student_model_type + ) + self.generation_config = SimpleNamespace(eos_token_id=2) + self.name_or_path = "student" + + class DummyTeacherModel: + def __init__(self): + self.config = SimpleNamespace(_name_or_path="teacher", vision_config=True, model_type=teacher_model_type) + self.resized_to = None + + def resize_token_embeddings(self, vocab_size): + self.resized_to = vocab_size + + return DummyStudentModel(), DummyTeacherModel() + + +def _make_vlm_trainer_args(use_vllm=False): + """Helper to create minimal GOLDTrainer args for VLM tests.""" + return SimpleNamespace( + model_init_kwargs=None, + max_length=128, + use_liger_kernel=False, + teacher_model_init_kwargs=None, + use_uld_loss=False, + teacher_tokenizer_name_or_path=None, + teacher_model_revision=None, + disable_dropout=False, + lmbda=0.5, + beta=0.5, + temperature=1.0, + top_p=1.0, + seq_kd=False, + num_generations=1, + use_transformers_paged=False, + max_completion_length=16, + top_k=0, + log_completions=False, + log_completions_steps=100, + wandb_log_unique_prompts=True, + num_completions_to_print=None, + per_device_train_batch_size=1, + gradient_accumulation_steps=1, + use_vllm=use_vllm, + vllm_mode="colocate", + vllm_structured_outputs_regex=None, + vllm_server_base_url=None, + vllm_server_host="0.0.0.0", + vllm_server_port=8001, + vllm_group_port=51216, + vllm_server_timeout=240.0, + vllm_tensor_parallel_size=1, + vllm_gpu_memory_utilization=0.2, + vllm_max_model_length=None, + vllm_enable_sleep_mode=False, + vllm_model_impl="vllm", + vllm_sync_frequency=1, + # ULD-specific defaults (needed when use_uld_loss=True) + uld_crossentropy_weight=0.5, + uld_distillation_weight=0.5, + uld_student_temperature=1.0, + uld_teacher_temperature=1.0, + uld_skip_student_eos=False, + uld_skip_teacher_eos=False, + use_extended_uld=False, + ) + + +def test_cross_architecture_vlm_without_uld_raises_error(monkeypatch): + """When student and teacher have different model_type and use_uld_loss=False, GOLDTrainer should raise + a ValueError telling the user to enable ULD loss.""" + + def fake_sft_init( + self, + model, + args=None, + data_collator=None, + train_dataset=None, + eval_dataset=None, + processing_class=None, + compute_metrics=None, + callbacks=None, + optimizers=None, + preprocess_logits_for_metrics=None, + peft_config=None, + ): + self.data_collator = data_collator + self.model = model + self.args = args + self.processing_class = processing_class + self.accelerator = SimpleNamespace( + device=torch.device("cpu"), + num_processes=1, + prepare_model=lambda module, evaluation_mode=True: module, + ) + self.is_deepspeed_enabled = False + self.is_fsdp_enabled = False + + monkeypatch.setattr(gold_trainer_module.SFTTrainer, "__init__", fake_sft_init) + + processor = AutoProcessor.from_pretrained("HuggingFaceTB/SmolVLM-256M-Instruct") + if processor.tokenizer.pad_token is None: + processor.tokenizer.pad_token = processor.tokenizer.eos_token + + sentinel_processor = SimpleNamespace(_is_sentinel=True) + real_auto_processor_from_pretrained = AutoProcessor.from_pretrained + + def patched_auto_processor(name, **kwargs): + if name == "teacher": + return sentinel_processor + return real_auto_processor_from_pretrained(name, **kwargs) + + monkeypatch.setattr(gold_trainer_module.AutoProcessor, "from_pretrained", staticmethod(patched_auto_processor)) + + vision_dataset = Dataset.from_dict({"messages": [["dummy"]], "image": ["fake_image"]}) + student, teacher = _make_dummy_vlm_models("smolvlm", "qwen2_5_vl") + args = _make_vlm_trainer_args() # use_uld_loss=False by default + + with pytest.raises(ValueError, match="Cross-architecture VLM distillation.*use_uld_loss=True"): + GOLDTrainer( + model=student, + teacher_model=teacher, + args=args, + train_dataset=vision_dataset, + processing_class=processor, + ) + + +def test_cross_architecture_vlm_with_uld_sets_teacher_processor(monkeypatch): + """When student and teacher have different model_type and use_uld_loss=True, GOLDTrainer should store + a separate _teacher_processor and emit a warning.""" + + def fake_sft_init( + self, + model, + args=None, + data_collator=None, + train_dataset=None, + eval_dataset=None, + processing_class=None, + compute_metrics=None, + callbacks=None, + optimizers=None, + preprocess_logits_for_metrics=None, + peft_config=None, + ): + self.data_collator = data_collator + self.model = model + self.args = args + self.processing_class = processing_class + self.accelerator = SimpleNamespace( + device=torch.device("cpu"), + num_processes=1, + prepare_model=lambda module, evaluation_mode=True: module, + ) + self.is_deepspeed_enabled = False + self.is_fsdp_enabled = False + + monkeypatch.setattr(gold_trainer_module.SFTTrainer, "__init__", fake_sft_init) + + processor = AutoProcessor.from_pretrained("HuggingFaceTB/SmolVLM-256M-Instruct") + if processor.tokenizer.pad_token is None: + processor.tokenizer.pad_token = processor.tokenizer.eos_token + + sentinel_processor = SimpleNamespace(_is_sentinel=True) + real_auto_processor_from_pretrained = AutoProcessor.from_pretrained + + def patched_auto_processor(name, **kwargs): + if name == "teacher": + return sentinel_processor + return real_auto_processor_from_pretrained(name, **kwargs) + + monkeypatch.setattr(gold_trainer_module.AutoProcessor, "from_pretrained", staticmethod(patched_auto_processor)) + + # Monkeypatch AutoTokenizer.from_pretrained for ULD teacher tokenizer loading + sentinel_tokenizer = SimpleNamespace(pad_token="", eos_token="") + real_auto_tokenizer_from_pretrained = AutoTokenizer.from_pretrained + + def patched_auto_tokenizer(name, **kwargs): + if name == "teacher": + return sentinel_tokenizer + return real_auto_tokenizer_from_pretrained(name, **kwargs) + + monkeypatch.setattr(gold_trainer_module.AutoTokenizer, "from_pretrained", staticmethod(patched_auto_tokenizer)) + + vision_dataset = Dataset.from_dict({"messages": [["dummy"]], "image": ["fake_image"]}) + student, teacher = _make_dummy_vlm_models("smolvlm", "qwen2_5_vl") + args = _make_vlm_trainer_args() + args.use_uld_loss = True + args.teacher_tokenizer_name_or_path = "teacher" + + import warnings + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + trainer = GOLDTrainer( + model=student, + teacher_model=teacher, + args=args, + train_dataset=vision_dataset, + processing_class=processor, + ) + + # _teacher_processor should be set for cross-architecture + assert trainer._teacher_processor is not None + assert trainer._teacher_processor is sentinel_processor + + # A cross-architecture warning should have been emitted + cross_arch_warnings = [w for w in caught if "Cross-architecture VLM distillation" in str(w.message)] + assert len(cross_arch_warnings) == 1 + assert "smolvlm" in str(cross_arch_warnings[0].message) + assert "qwen2_5_vl" in str(cross_arch_warnings[0].message) + + # Identity collator and VLM collator should still be set + assert trainer.data_collator is identity + assert trainer._vlm_collator is not None + + +def test_same_architecture_vlm_no_teacher_processor(monkeypatch): + """When student and teacher have the same model_type, GOLDTrainer should NOT store a _teacher_processor + (zero overhead -- both models share the same forward_kwargs).""" + + def fake_sft_init( + self, + model, + args=None, + data_collator=None, + train_dataset=None, + eval_dataset=None, + processing_class=None, + compute_metrics=None, + callbacks=None, + optimizers=None, + preprocess_logits_for_metrics=None, + peft_config=None, + ): + self.data_collator = data_collator + self.model = model + self.args = args + self.processing_class = processing_class + self.accelerator = SimpleNamespace( + device=torch.device("cpu"), + num_processes=1, + prepare_model=lambda module, evaluation_mode=True: module, + ) + self.is_deepspeed_enabled = False + self.is_fsdp_enabled = False + + monkeypatch.setattr(gold_trainer_module.SFTTrainer, "__init__", fake_sft_init) + + processor = AutoProcessor.from_pretrained("HuggingFaceTB/SmolVLM-256M-Instruct") + if processor.tokenizer.pad_token is None: + processor.tokenizer.pad_token = processor.tokenizer.eos_token + + vision_dataset = Dataset.from_dict({"messages": [["dummy"]], "image": ["fake_image"]}) + student, teacher = _make_dummy_vlm_models("smolvlm", "smolvlm") + args = _make_vlm_trainer_args() + + import warnings + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + trainer = GOLDTrainer( + model=student, + teacher_model=teacher, + args=args, + train_dataset=vision_dataset, + processing_class=processor, + ) + + # _teacher_processor should be None for same architecture (zero overhead) + assert trainer._teacher_processor is None + + # No cross-architecture warning should have been emitted + cross_arch_warnings = [w for w in caught if "Cross-architecture VLM distillation" in str(w.message)] + assert len(cross_arch_warnings) == 0 + + # Identity collator and VLM collator should still be set + assert trainer.data_collator is identity + assert trainer._vlm_collator is not None diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index 933189a0970..973e46dccff 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -29,7 +29,7 @@ from accelerate.utils import DistributedType, broadcast_object_list, gather_object from datasets import Dataset, IterableDataset from torch.utils.data import DataLoader -from transformers import AutoProcessor, AutoTokenizer, TrainerCallback +from transformers import AutoConfig, AutoProcessor, AutoTokenizer, TrainerCallback from transformers.data.data_collator import DataCollator from transformers.feature_extraction_utils import FeatureExtractionMixin from transformers.generation.configuration_utils import GenerationConfig @@ -806,14 +806,25 @@ def __init__( # VLM distillation: only VLM-to-VLM is supported. Both student and teacher must be # VLMs so that both receive images and multimodal inputs. + self._teacher_processor = None if self._is_vlm and isinstance(teacher_model, str): - # Teacher not yet instantiated + # Teacher not yet instantiated -- validate it's a VLM teacher_proc = AutoProcessor.from_pretrained(teacher_model) if not isinstance(teacher_proc, ProcessorMixin): raise ValueError( "VLM distillation requires both student and teacher to be vision-language models. " "The student has a `ProcessorMixin` but the teacher does not." ) + # Check for cross-architecture VLM distillation + student_model_type = model.config.model_type if not isinstance(model, str) else None + teacher_model_type = AutoConfig.from_pretrained(teacher_model).model_type + if student_model_type and teacher_model_type != student_model_type: + warnings.warn( + f"Cross-architecture VLM distillation detected: student is '{student_model_type}', " + f"teacher is '{teacher_model_type}'. Images will be processed separately through each " + "model's processor, which may increase memory usage and computation time." + ) + self._teacher_processor = teacher_proc elif self._is_vlm and not isinstance(teacher_model, str): # Teacher already instantiated — check if it looks like a VLM by checking for a vision config if not hasattr(teacher_model, "config") or not hasattr(teacher_model.config, "vision_config"): @@ -822,6 +833,23 @@ def __init__( "The student has a `ProcessorMixin` but the teacher model does not appear to be a VLM " "(missing `vision_config`)." ) + # Check for cross-architecture VLM distillation + student_model_type = model.config.model_type if not isinstance(model, str) else None + teacher_model_type = teacher_model.config.model_type + if student_model_type and teacher_model_type != student_model_type: + warnings.warn( + f"Cross-architecture VLM distillation detected: student is '{student_model_type}', " + f"teacher is '{teacher_model_type}'. Images will be processed separately through each " + "model's processor, which may increase memory usage and computation time." + ) + self._teacher_processor = AutoProcessor.from_pretrained(teacher_model.config._name_or_path) + if self._teacher_processor is not None and not args.use_uld_loss: + raise ValueError( + "Cross-architecture VLM distillation (student and teacher have different `model_type`) is not " + "supported with the standard JSD loss because the models require different image token formats " + "and tokenizers. Please set `use_uld_loss=True` in your GOLDConfig to enable cross-tokenizer " + "alignment via ULD loss." + ) self._is_vision_dataset = "image" in dataset_sample or "images" in dataset_sample if self._is_vision_dataset and not self._is_vlm: raise ValueError( @@ -830,22 +858,17 @@ def __init__( ) # Respect a user-provided data_collator; otherwise, pick the right collator based on modality. - # For VLMs with lmbda > 0, use identity collator to preserver raw PIL images in the dataloader. - # Because vLLM requires raw images, not processed pixel_values tensors. + # For VLMs, always use identity collator to preserve raw PIL images in the dataloader. + # Raw images are needed for: (1) vLLM generation, (2) cross-architecture teacher processing. # A separate _vlm_collator is stored for on-the-fly collation inside _fill_buffer. self._vlm_collator = None if data_collator is None: - if self._is_vision_dataset and args.lmbda > 0: + if self._is_vision_dataset: self._vlm_collator = DataCollatorForVisionLanguageChatML( processor=processing_class, max_length=args.max_length, ) data_collator = identity - elif self._is_vision_dataset: - data_collator = DataCollatorForVisionLanguageChatML( - processor=processing_class, - max_length=args.max_length, - ) else: data_collator = DataCollatorForChatML(tokenizer=processing_class, max_length=args.max_length) @@ -879,6 +902,8 @@ def __init__( if args.use_uld_loss and args.teacher_tokenizer_name_or_path is None: if isinstance(teacher_model, str): args.teacher_tokenizer_name_or_path = teacher_model + elif hasattr(teacher_model, "config") and getattr(teacher_model.config, "_name_or_path", None): + args.teacher_tokenizer_name_or_path = teacher_model.config._name_or_path else: raise ValueError( "`teacher_tokenizer_name_or_path` must be set when using ULD loss with a pre-instantiated teacher model." @@ -1229,6 +1254,12 @@ def _fill_buffer(self, generation_batch: dict[str, torch.Tensor | Any] | list[di k: v.to(self.accelerator.device) if isinstance(v, torch.Tensor) else v for k, v in slice_inputs.items() } + # Preserve raw PIL images and prompts for cross-architecture teacher processing + if self._teacher_processor is not None: + slice_inputs["_raw_images"] = [ + ex.get("images") or ([ex["image"]] if "image" in ex else None) for ex in raw_slices[i] + ] + slice_inputs["_raw_prompts"] = [ex.get("prompt") for ex in raw_slices[i]] else: slice_inputs = slices[i] @@ -1404,6 +1435,9 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in updated_slice["labels"] = new_labels updated_slice["original_prompt_text"] = prompt_texts updated_slice["original_completion_text"] = completion_texts + if self._teacher_processor is not None: + updated_slice["_raw_images"] = images + updated_slice["_raw_prompts"] = prompts self._buffered_inputs[slice_idx] = updated_slice self._buffered_text_logs[slice_idx] = (prompt_texts, completion_texts) @@ -1447,6 +1481,9 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in collated = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in collated.items()} collated["original_prompt_text"] = prompts_text_with_special collated["original_completion_text"] = completion_texts + if self._teacher_processor is not None: + collated["_raw_images"] = images + collated["_raw_prompts"] = prompts self._buffered_inputs[slice_idx] = collated self._buffered_text_logs[slice_idx] = (prompts_text, completion_texts) @@ -1922,8 +1959,11 @@ def generalized_jsd_loss( ) def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): - # Extract multimodal fields for VLM forward passes - forward_kwargs = {k: inputs[k] for k in self._MULTIMODAL_KEYS if k in inputs} + # Extract multimodal fields for student forward passes + student_forward_kwargs = {k: inputs[k] for k in self._MULTIMODAL_KEYS if k in inputs} + # For same-architecture teacher reuses student vision tensors. + # For cross-architecture VLMs, this gets overridden in the ULD branch below. + teacher_forward_kwargs = student_forward_kwargs if self.use_uld_loss and self.teacher_tokenizer is not None: if "original_prompt_text" in inputs and "original_completion_text" in inputs: @@ -1943,16 +1983,60 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N full.replace(prompt, "", 1) for full, prompt in zip(full_texts, prompt_texts, strict=True) ] - ( - teacher_input_ids, - teacher_labels, - teacher_attention_mask, - teacher_prompt_length, - ) = build_teacher_inputs_from_texts( - self.teacher_tokenizer, - prompt_texts, - completion_texts, - ) + # For cross-architecture VLMs, build teacher inputs with image placeholders by processing + # prompts through the teacher's processor with raw images, then appending completions. + if self._teacher_processor is not None and "_raw_images" in inputs: + raw_images = inputs["_raw_images"] + raw_prompts = inputs["_raw_prompts"] + # Apply teacher's chat template to get prompt text with correct image placeholders + teacher_prompt_texts = self._teacher_processor.apply_chat_template( + raw_prompts, tokenize=False, add_generation_prompt=True + ) + # Build full text (prompt + completion) and process in one call so all tensors + # (input_ids, attention_mask, mm_token_type_ids, pixel_values, ...) are aligned. + teacher_full_texts = [p + c for p, c in zip(teacher_prompt_texts, completion_texts, strict=True)] + teacher_full_processed = self._teacher_processor( + images=raw_images, + text=teacher_full_texts, + padding=True, + return_tensors="pt", + ) + teacher_input_ids = teacher_full_processed["input_ids"] + teacher_attention_mask = teacher_full_processed["attention_mask"] + # Determine prompt lengths after image token expansion to build labels. + # Derive prompt lengths from total sequence length minus completion length. + # Completions are pure text (no images), so the tokenizer gives exact counts. + # This avoids a second image-processing pass through the teacher processor. + teacher_completion_token_lengths = [ + len(self._teacher_processor.tokenizer(ct, add_special_tokens=False)["input_ids"]) + for ct in completion_texts + ] + total_lengths = teacher_attention_mask.sum(dim=1) + teacher_prompt_token_lengths = [ + int(total_lengths[i].item()) - cl for i, cl in enumerate(teacher_completion_token_lengths) + ] + teacher_labels = teacher_input_ids.clone() + teacher_labels[teacher_attention_mask == 0] = -100 + for i, pl in enumerate(teacher_prompt_token_lengths): + teacher_labels[i, :pl] = -100 + teacher_prompt_length = max(teacher_prompt_token_lengths) + # Override teacher_forward_kwargs with all multimodal keys from teacher processing + teacher_forward_kwargs = { + k: teacher_full_processed[k].to(self.accelerator.device) + for k in self._MULTIMODAL_KEYS + if k in teacher_full_processed + } + else: + ( + teacher_input_ids, + teacher_labels, + teacher_attention_mask, + teacher_prompt_length, + ) = build_teacher_inputs_from_texts( + self.teacher_tokenizer, + prompt_texts, + completion_texts, + ) teacher_input_ids = teacher_input_ids.to(self.accelerator.device) teacher_labels = teacher_labels.to(self.accelerator.device) @@ -1962,7 +2046,7 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"], use_cache=False, - **forward_kwargs, + **student_forward_kwargs, ) self.teacher_model.eval() @@ -1970,7 +2054,7 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N outputs_teacher = self.teacher_model( input_ids=teacher_input_ids, attention_mask=teacher_attention_mask, - **forward_kwargs, + **teacher_forward_kwargs, ) # These are not used for ULD loss but are needed if JSD loss were to be used in this branch @@ -1993,7 +2077,7 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"], use_cache=False, - **forward_kwargs, + **student_forward_kwargs, ) self.teacher_model.eval() @@ -2009,7 +2093,7 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"], use_cache=False, - **forward_kwargs, + **teacher_forward_kwargs, ) student_hidden = student_outputs.last_hidden_state[:, :-1] @@ -2044,7 +2128,7 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N outputs_student = model( input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"], - **forward_kwargs, + **student_forward_kwargs, ) self.teacher_model.eval() @@ -2052,7 +2136,7 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N outputs_teacher = self.teacher_model( input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"], - **forward_kwargs, + **teacher_forward_kwargs, ) prompt_lengths = inputs["prompts"].shape[1] From 9a1f345daf696f579aa08809bcf86aa25b582728 Mon Sep 17 00:00:00 2001 From: Strongich Date: Mon, 6 Apr 2026 20:01:38 +0200 Subject: [PATCH 04/39] fix collator mutation bug and reject Liger kernel for VLMs --- trl/experimental/gold/gold_trainer.py | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index 973e46dccff..d88d354ea38 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -875,6 +875,12 @@ def __init__( # Liger fused GKD loss (JSD) self.use_liger_gkd_loss = False if args.use_liger_kernel: + if self._is_vlm: + raise ValueError( + "Liger fused GKD loss is not supported with VLMs. The fused kernel operates on base decoder " + "hidden states, which is incompatible with VLM multimodal inputs (pixel_values, etc.). " + "Please set `use_liger_kernel=False`." + ) self.liger_jsd_loss = LigerFusedLinearJSDLoss( beta=args.beta, ignore_index=-100, @@ -1248,6 +1254,15 @@ def _fill_buffer(self, generation_batch: dict[str, torch.Tensor | Any] | list[di for i, flag in enumerate(on_policy_flags): if not flag: if self._vlm_collator is not None: + # Extract raw images and prompts BEFORE collation, since the collator + # mutates examples in place (pops "image", overwrites "prompt"). + raw_images = None + raw_prompts = None + if self._teacher_processor is not None: + raw_images = [ + ex.get("images") or ([ex["image"]] if "image" in ex else None) for ex in raw_slices[i] + ] + raw_prompts = [ex.get("prompt") for ex in raw_slices[i]] # Collate raw examples on-the-fly for off-policy slices slice_inputs = self._vlm_collator(raw_slices[i]) slice_inputs = { @@ -1256,10 +1271,8 @@ def _fill_buffer(self, generation_batch: dict[str, torch.Tensor | Any] | list[di } # Preserve raw PIL images and prompts for cross-architecture teacher processing if self._teacher_processor is not None: - slice_inputs["_raw_images"] = [ - ex.get("images") or ([ex["image"]] if "image" in ex else None) for ex in raw_slices[i] - ] - slice_inputs["_raw_prompts"] = [ex.get("prompt") for ex in raw_slices[i]] + slice_inputs["_raw_images"] = raw_images + slice_inputs["_raw_prompts"] = raw_prompts else: slice_inputs = slices[i] From 606d68d9ec6b54bf3c11ef2b58c5543d35460f61 Mon Sep 17 00:00:00 2001 From: Strongich Date: Mon, 6 Apr 2026 20:42:20 +0200 Subject: [PATCH 05/39] pass tokenizer instead of processing_class to the ULDLoss --- trl/experimental/gold/gold_trainer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index d88d354ea38..4205231e502 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -986,7 +986,7 @@ def __init__( if self.use_uld_loss: self.uld_loss_fn = ULDLoss( config=args, - student_tokenizer=processing_class, + student_tokenizer=tokenizer, teacher_tokenizer=self.teacher_tokenizer, device=self.accelerator.device, ) From 159e7a3378e8279f883eb71878537cdb5a802fb8 Mon Sep 17 00:00:00 2001 From: Strongich Date: Tue, 7 Apr 2026 12:22:01 +0200 Subject: [PATCH 06/39] fix prompt_lengths split to use min instead of max after flush_left --- trl/experimental/gold/gold_trainer.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index 4205231e502..9289367b27b 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -226,7 +226,7 @@ def build_teacher_inputs_from_texts( last_idx = valid.nonzero(as_tuple=True)[0][-1] teacher_attention_mask[row, last_idx + 1 :] = False - teacher_prompt_length = max(prompt_lengths) if prompt_lengths else 0 + teacher_prompt_length = min(prompt_lengths) if prompt_lengths else 0 return teacher_input_ids, teacher_labels, teacher_attention_mask, teacher_prompt_length @@ -2032,7 +2032,7 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N teacher_labels[teacher_attention_mask == 0] = -100 for i, pl in enumerate(teacher_prompt_token_lengths): teacher_labels[i, :pl] = -100 - teacher_prompt_length = max(teacher_prompt_token_lengths) + teacher_prompt_length = min(teacher_prompt_token_lengths) # Override teacher_forward_kwargs with all multimodal keys from teacher processing teacher_forward_kwargs = { k: teacher_full_processed[k].to(self.accelerator.device) @@ -2071,7 +2071,12 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N ) # These are not used for ULD loss but are needed if JSD loss were to be used in this branch - student_prompt_length = inputs["prompts"].shape[1] + # For VLMs, prompts are left-padded but input_ids are flushed left, so prompts.shape[1] + # would overcount. Derive prompt length from the flushed labels instead. + if self._is_vlm: + student_prompt_length = (inputs["labels"] != -100).long().argmax(dim=1).min().item() + else: + student_prompt_length = inputs["prompts"].shape[1] shifted_student_logits = outputs_student.logits[:, student_prompt_length - 1 : -1, :] shifted_teacher_logits = outputs_teacher.logits[:, teacher_prompt_length - 1 : -1, :] shifted_labels = inputs["labels"][:, student_prompt_length:] @@ -2151,7 +2156,6 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N attention_mask=inputs["attention_mask"], **teacher_forward_kwargs, ) - prompt_lengths = inputs["prompts"].shape[1] shifted_student_logits = outputs_student.logits[:, prompt_lengths - 1 : -1, :] shifted_teacher_logits = outputs_teacher.logits[:, prompt_lengths - 1 : -1, :] From bd820ad8c32510f7e9d4ab5b1a83a955178f3130 Mon Sep 17 00:00:00 2001 From: Strongich Date: Tue, 7 Apr 2026 14:06:16 +0200 Subject: [PATCH 07/39] fix prompt length split for JSD loss --- trl/experimental/gold/gold_trainer.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index 9289367b27b..2ee2127efe1 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -2156,7 +2156,12 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N attention_mask=inputs["attention_mask"], **teacher_forward_kwargs, ) - prompt_lengths = inputs["prompts"].shape[1] + # Using the same prompt_lengths for teacher and student, since JSD can only be + # used with same-family VLMs (shared tokenizer). + if self._is_vlm: + prompt_lengths = (inputs["labels"] != -100).long().argmax(dim=1).min().item() + else: + prompt_lengths = inputs["prompts"].shape[1] shifted_student_logits = outputs_student.logits[:, prompt_lengths - 1 : -1, :] shifted_teacher_logits = outputs_teacher.logits[:, prompt_lengths - 1 : -1, :] shifted_labels = inputs["labels"][:, prompt_lengths:] From 571099ec37563621ab2290d313f80b2ffe8399e0 Mon Sep 17 00:00:00 2001 From: Strongich Date: Tue, 7 Apr 2026 16:36:07 +0200 Subject: [PATCH 08/39] batch VLM vLLM generation across slices & fix VLM dataset columns stripped --- trl/experimental/gold/gold_config.py | 9 +++ trl/experimental/gold/gold_trainer.py | 111 ++++++++++++++++++-------- 2 files changed, 88 insertions(+), 32 deletions(-) diff --git a/trl/experimental/gold/gold_config.py b/trl/experimental/gold/gold_config.py index 1af9eeae332..09f6689292d 100644 --- a/trl/experimental/gold/gold_config.py +++ b/trl/experimental/gold/gold_config.py @@ -119,6 +119,15 @@ class GOLDConfig(SFTConfig): default=1e-7, metadata={"help": "The initial learning rate for AdamW."}, ) + # The default value remove_unused_columns is overwritten from the parent class, because in GOLD we usually rely on + # additional columns to compute the loss + remove_unused_columns: bool | None = field( + default=False, + metadata={ + "help": "Whether to only keep the columns 'prompt' and 'completion' in the dataset. If you use a custom " + "dataset that requires additional columns, you should keep this to `False`." + }, + ) # GOLD-specific parameters temperature: float = field( diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index 2ee2127efe1..b23df7c3590 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -1064,6 +1064,8 @@ def __init__( def _set_signature_columns_if_needed(self): super()._set_signature_columns_if_needed() required_columns = [ + "prompt", + "completion", "prompts", "prompt_attention_mask", "messages", @@ -1378,6 +1380,14 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in """On-policy generation from raw VLM examples, preserving PIL images for vLLM.""" device = self.accelerator.device + # Phase 1: Collect prompts, images, and raw examples across all on-policy slices + all_prompt_ids = [] + all_images = [] + all_prompts = [] # prepared multimodal messages + all_raw_examples = [] + local_slice_indices = [] + slice_raw_data = {} # per-slice raw data for non-vLLM path + for slice_idx in on_policy_indices: raw_examples = raw_slices[slice_idx] @@ -1427,11 +1437,22 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in for ids, mask in zip(tokenized["input_ids"], tokenized["attention_mask"], strict=True) ] - prompts_text = self.processing_class.batch_decode(prompt_ids_list, skip_special_tokens=True) - prompts_text_with_special = self.processing_class.batch_decode(prompt_ids_list, skip_special_tokens=False) + slice_raw_data[slice_idx] = (raw_examples, images, prompts, prompt_ids_list) + + for i, example in enumerate(raw_examples): + all_prompt_ids.append(prompt_ids_list[i]) + all_images.append(images[i] if images is not None else None) + all_prompts.append(prompts[i]) + all_raw_examples.append(example) + local_slice_indices.append(slice_idx) + + all_prompts_text = self.processing_class.batch_decode(all_prompt_ids, skip_special_tokens=True) + all_prompts_text_with_special = self.processing_class.batch_decode(all_prompt_ids, skip_special_tokens=False) - if not self.use_vllm: - # Non-vLLM path: collate raw examples to get pixel_values, then generate + if not self.use_vllm: + # Non-vLLM path: generate per-slice using model.generate + for slice_idx in on_policy_indices: + raw_examples, images, prompts, _ = slice_raw_data[slice_idx] collated = self._vlm_collator(raw_examples) collated = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in collated.items()} with unwrap_model_for_generation( @@ -1454,37 +1475,62 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in self._buffered_inputs[slice_idx] = updated_slice self._buffered_text_logs[slice_idx] = (prompt_texts, completion_texts) - continue + return - # vLLM path: pass raw PIL images to vLLM - if ( - self.state.global_step != self._last_vllm_sync_step - and self.state.global_step >= self._last_vllm_sync_step + self.vllm_sync_frequency - ): - self.vllm_generation.sync_weights() - self._last_vllm_sync_step = self.state.global_step + # vLLM path: one batched generate call across all slices + if ( + self.state.global_step != self._last_vllm_sync_step + and self.state.global_step >= self._last_vllm_sync_step + self.vllm_sync_frequency + ): + self.vllm_generation.sync_weights() + self._last_vllm_sync_step = self.state.global_step + + # Pass None for images if all entries are None + generate_images = all_images if any(img is not None for img in all_images) else None + _, completion_ids, _, _ = self.vllm_generation.generate( + prompts=all_prompt_ids, + images=generate_images, + num_generations=self.num_generations, + ) - _, completion_ids, _, _ = self.vllm_generation.generate( - prompts=prompt_ids_list, - images=images, - num_generations=self.num_generations, + # Decode completions + max_completion_length = self.generation_config.max_new_tokens + all_completion_texts = [] + for comp_ids in completion_ids: + if len(comp_ids) > max_completion_length: + comp_ids = comp_ids[:max_completion_length] + all_completion_texts.append( + self.processing_class.decode(comp_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False) ) - # Decode completions and build synthetic examples for collation - max_completion_length = self.generation_config.max_new_tokens - completion_texts = [] - for comp_ids in completion_ids: - if len(comp_ids) > max_completion_length: - comp_ids = comp_ids[:max_completion_length] - completion_texts.append( - self.processing_class.decode( - comp_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False - ) - ) + # Redistribute completions to slices. With num_generations > 1, each prompt produces + # multiple completions, so we repeat each raw example/prompt/image to match. + slice_completions = {idx: [] for idx in on_policy_indices} + slice_raw = {idx: [] for idx in on_policy_indices} + slice_images = {idx: [] for idx in on_policy_indices} + slice_prompts = {idx: [] for idx in on_policy_indices} + slice_prompts_text = {idx: [] for idx in on_policy_indices} + slice_prompts_text_special = {idx: [] for idx in on_policy_indices} + + for i, slice_idx in enumerate(local_slice_indices): + for g in range(self.num_generations): + comp_idx = i * self.num_generations + g + slice_completions[slice_idx].append(all_completion_texts[comp_idx]) + slice_raw[slice_idx].append(all_raw_examples[i]) + slice_images[slice_idx].append(all_images[i]) + slice_prompts[slice_idx].append(all_prompts[i]) + slice_prompts_text[slice_idx].append(all_prompts_text[i]) + slice_prompts_text_special[slice_idx].append(all_prompts_text_with_special[i]) + + for slice_idx in on_policy_indices: + completion_texts = slice_completions[slice_idx] + raw_for_slice = slice_raw[slice_idx] + images_for_slice = slice_images[slice_idx] + prompts_for_slice = slice_prompts[slice_idx] # Build synthetic examples: original prompt + generated completion synthetic_examples = [] - for i, example in enumerate(raw_examples): + for i, example in enumerate(raw_for_slice): synthetic = dict(example) synthetic["completion"] = [{"role": "assistant", "content": completion_texts[i]}] synthetic_examples.append(synthetic) @@ -1492,14 +1538,15 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in # Collate synthetic examples to get pixel_values + properly tokenized input_ids/labels collated = self._vlm_collator(synthetic_examples) collated = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in collated.items()} - collated["original_prompt_text"] = prompts_text_with_special + collated["original_prompt_text"] = slice_prompts_text_special[slice_idx] collated["original_completion_text"] = completion_texts if self._teacher_processor is not None: - collated["_raw_images"] = images - collated["_raw_prompts"] = prompts + has_images = any(img is not None for img in images_for_slice) + collated["_raw_images"] = images_for_slice if has_images else None + collated["_raw_prompts"] = prompts_for_slice self._buffered_inputs[slice_idx] = collated - self._buffered_text_logs[slice_idx] = (prompts_text, completion_texts) + self._buffered_text_logs[slice_idx] = (slice_prompts_text[slice_idx], completion_texts) def _process_completions_to_buffer( self, From eaa258d742e62db56957b4dccf31420564c66b84 Mon Sep 17 00:00:00 2001 From: Strongich Date: Tue, 7 Apr 2026 22:16:13 +0200 Subject: [PATCH 09/39] fix VLM data pipeline --- trl/experimental/gold/gold_trainer.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index b23df7c3590..9491e004d5a 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -1264,7 +1264,12 @@ def _fill_buffer(self, generation_batch: dict[str, torch.Tensor | Any] | list[di raw_images = [ ex.get("images") or ([ex["image"]] if "image" in ex else None) for ex in raw_slices[i] ] - raw_prompts = [ex.get("prompt") for ex in raw_slices[i]] + raw_prompts = [ + prepare_multimodal_messages(ex["prompt"], images=imgs) + if imgs is not None + else ex.get("prompt") + for ex, imgs in zip(raw_slices[i], raw_images, strict=True) + ] # Collate raw examples on-the-fly for off-policy slices slice_inputs = self._vlm_collator(raw_slices[i]) slice_inputs = { @@ -1500,7 +1505,7 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in if len(comp_ids) > max_completion_length: comp_ids = comp_ids[:max_completion_length] all_completion_texts.append( - self.processing_class.decode(comp_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False) + self.processing_class.decode(comp_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False) ) # Redistribute completions to slices. With num_generations > 1, each prompt produces From 77f4ccf200ec7a62aec130d29c23f05b25c7b9fb Mon Sep 17 00:00:00 2001 From: Strongich Date: Wed, 8 Apr 2026 18:24:39 +0200 Subject: [PATCH 10/39] fix multimodal key shape missmatch for on-policy gen --- trl/experimental/gold/gold_trainer.py | 32 ++++ trl/experimental/utils.py | 253 ++++++++++++++++++++------ 2 files changed, 226 insertions(+), 59 deletions(-) diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index 9491e004d5a..ac2e574409c 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -1375,6 +1375,19 @@ def _generate_non_vllm_for_slices(self, slices: list[dict[str, torch.Tensor | An updated_slice["input_ids"] = new_input_ids updated_slice["attention_mask"] = new_attention_mask updated_slice["labels"] = new_labels + # Rebuild sequence-length-dependent keys to match new input_ids shape + new_seq_len = new_input_ids.shape[1] + prompt_seq_len = slice_inputs["prompts"].shape[1] + for k in ("token_type_ids", "mm_token_type_ids"): + if k in updated_slice: + prompt_part = updated_slice[k][:, :prompt_seq_len] + comp_part = torch.zeros( + new_input_ids.shape[0], + new_seq_len - prompt_seq_len, + dtype=updated_slice[k].dtype, + device=new_input_ids.device, + ) + updated_slice[k] = torch.cat([prompt_part, comp_part], dim=1) updated_slice["original_prompt_text"] = prompt_texts updated_slice["original_completion_text"] = completion_texts @@ -1472,6 +1485,19 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in updated_slice["input_ids"] = new_input_ids updated_slice["attention_mask"] = new_attention_mask updated_slice["labels"] = new_labels + # Rebuild sequence-length-dependent keys to match new input_ids shape + new_seq_len = new_input_ids.shape[1] + prompt_seq_len = collated["prompts"].shape[1] + for k in ("token_type_ids", "mm_token_type_ids"): + if k in updated_slice: + prompt_part = updated_slice[k][:, :prompt_seq_len] + comp_part = torch.zeros( + new_input_ids.shape[0], + new_seq_len - prompt_seq_len, + dtype=updated_slice[k].dtype, + device=new_input_ids.device, + ) + updated_slice[k] = torch.cat([prompt_part, comp_part], dim=1) updated_slice["original_prompt_text"] = prompt_texts updated_slice["original_completion_text"] = completion_texts if self._teacher_processor is not None: @@ -2298,6 +2324,12 @@ def generate_on_policy_outputs(self, model, inputs, generation_config, pad_token generated_tokens = torch.stack([torch.tensor(ids, device=model.device) for ids in completion_ids]) else: generate_kwargs = {k: inputs[k] for k in self._MULTIMODAL_KEYS if k in inputs} + # Slice sequence-length-dependent keys to prompt-only length (e.g. token_type_ids for Gemma, + # mm_token_type_ids for ERNIE-VL) since model.generate receives prompt-only input_ids + prompt_seq_len = inputs["prompts"].shape[1] + for k in ("token_type_ids", "mm_token_type_ids"): + if k in generate_kwargs: + generate_kwargs[k] = generate_kwargs[k][:, :prompt_seq_len] generated_outputs = model.generate( input_ids=inputs["prompts"], attention_mask=inputs.get("prompt_attention_mask", None), diff --git a/trl/experimental/utils.py b/trl/experimental/utils.py index 70a41580a74..5526439717a 100644 --- a/trl/experimental/utils.py +++ b/trl/experimental/utils.py @@ -73,7 +73,9 @@ def __call__(self, features: list[dict[str, Any]]) -> dict[str, Any]: # first, pad everything to the same length padded_batch = {} for k in features[0].keys(): - if k.endswith(("_input_ids", "_attention_mask", "_labels", "_pixel_values")): + if k.endswith( + ("_input_ids", "_attention_mask", "_labels", "_pixel_values") + ): if self.is_encoder_decoder: to_pad = [torch.LongTensor(ex[k]) for ex in features] @@ -87,11 +89,15 @@ def __call__(self, features: list[dict[str, Any]]) -> dict[str, Any]: padding_value = self.pad_token_id elif k.endswith("_attention_mask"): padding_value = 0 - elif k.startswith(("chosen", "rejected", "completion")) or ("decoder" in k): + elif k.startswith(("chosen", "rejected", "completion")) or ( + "decoder" in k + ): padding_value = -100 else: raise ValueError(f"Unexpected key in batch '{k}'") - padded_batch[k] = pad_sequence(to_pad, batch_first=True, padding_value=padding_value) + padded_batch[k] = pad_sequence( + to_pad, batch_first=True, padding_value=padding_value + ) else: # Set padding value based on the key if k.endswith("_input_ids"): @@ -119,13 +125,17 @@ def __call__(self, features: list[dict[str, Any]]) -> dict[str, Any]: # Set the dtype if k.endswith("_pixel_values"): - dtype = torch.float32 # will be downcasted if necessary by the Trainer + dtype = ( + torch.float32 + ) # will be downcasted if necessary by the Trainer else: dtype = torch.int64 # Convert to tensor and pad to_pad = [torch.tensor(ex[k], dtype=dtype) for ex in features] - padded_batch[k] = pad(to_pad, padding_value=padding_value, padding_side=padding_side) + padded_batch[k] = pad( + to_pad, padding_value=padding_value, padding_side=padding_side + ) elif k.endswith("_logps"): # the cached reference model logprobs padded_batch[k] = torch.tensor([ex[k] for ex in features]) @@ -149,7 +159,9 @@ class DataCollatorForChatML: def __post_init__(self): if self.tokenizer.pad_token_id is None: - raise ValueError("The tokenizer does not have a pad token. Please set `pad_token_id` in the tokenizer.") + raise ValueError( + "The tokenizer does not have a pad token. Please set `pad_token_id` in the tokenizer." + ) if self.max_length is None: # set a sensible default self.max_length = min(self.tokenizer.model_max_length, 1024) @@ -189,7 +201,11 @@ def __call__(self, examples: list[dict[str, Any]]) -> dict[str, torch.Tensor]: if offsets is not None: prompt_char_len = len(formatted_prompt) completion_start_idx_full = next( - (idx for idx, (start, _) in enumerate(offsets) if start >= prompt_char_len), + ( + idx + for idx, (start, _) in enumerate(offsets) + if start >= prompt_char_len + ), len(message_input_ids_full), ) else: @@ -203,16 +219,25 @@ def __call__(self, examples: list[dict[str, Any]]) -> dict[str, torch.Tensor]: completion_start_idx_full = len(tokenized_prompt_full["input_ids"]) prompt_tokens_full = message_input_ids_full[:completion_start_idx_full] - completion_input_ids_full = message_input_ids_full[completion_start_idx_full:] - - if self.max_length is not None and len(message_input_ids_full) > self.max_length: + completion_input_ids_full = message_input_ids_full[ + completion_start_idx_full: + ] + + if ( + self.max_length is not None + and len(message_input_ids_full) > self.max_length + ): completion_ids = completion_input_ids_full if len(completion_ids) >= self.max_length: completion_ids = completion_ids[-self.max_length :] prompt_ids = [] else: max_prompt_tokens = self.max_length - len(completion_ids) - prompt_ids = prompt_tokens_full[-max_prompt_tokens:] if max_prompt_tokens > 0 else [] + prompt_ids = ( + prompt_tokens_full[-max_prompt_tokens:] + if max_prompt_tokens > 0 + else [] + ) message_input_ids = prompt_ids + completion_ids else: message_input_ids = message_input_ids_full @@ -249,16 +274,30 @@ def __call__(self, examples: list[dict[str, Any]]) -> dict[str, torch.Tensor]: # convert to list of tensors and pad input_ids = [torch.tensor(ids, dtype=torch.long) for ids in input_ids] - attention_mask = [torch.tensor(mask, dtype=torch.long) for mask in attention_mask] + attention_mask = [ + torch.tensor(mask, dtype=torch.long) for mask in attention_mask + ] labels = [torch.tensor(label, dtype=torch.long) for label in labels] - input_ids = pad(input_ids, padding_side="left", padding_value=self.tokenizer.pad_token_id) + input_ids = pad( + input_ids, padding_side="left", padding_value=self.tokenizer.pad_token_id + ) attention_mask = pad(attention_mask, padding_side="left", padding_value=0) labels = pad(labels, padding_side="left", padding_value=self.ignore_index) - prompts_input_ids = [torch.tensor(ids, dtype=torch.long) for ids in prompts_input_ids] - prompt_attention_mask = [torch.tensor(mask, dtype=torch.long) for mask in prompt_attention_mask] - prompts_input_ids = pad(prompts_input_ids, padding_side="left", padding_value=self.tokenizer.pad_token_id) - prompt_attention_mask = pad(prompt_attention_mask, padding_side="left", padding_value=0) + prompts_input_ids = [ + torch.tensor(ids, dtype=torch.long) for ids in prompts_input_ids + ] + prompt_attention_mask = [ + torch.tensor(mask, dtype=torch.long) for mask in prompt_attention_mask + ] + prompts_input_ids = pad( + prompts_input_ids, + padding_side="left", + padding_value=self.tokenizer.pad_token_id, + ) + prompt_attention_mask = pad( + prompt_attention_mask, padding_side="left", padding_value=0 + ) return { "input_ids": input_ids, @@ -323,9 +362,15 @@ def torch_call(self, examples: list[dict[str, Any]]) -> dict[str, Any]: # Apply chat template for conversational data if is_conversational(examples[0]): for example in examples: - example["prompt"] = prepare_multimodal_messages(example["prompt"], images=example["images"]) - example["completion"] = prepare_multimodal_messages(example["completion"]) - examples = [apply_chat_template(example, self.processor) for example in examples] + example["prompt"] = prepare_multimodal_messages( + example["prompt"], images=example["images"] + ) + example["completion"] = prepare_multimodal_messages( + example["completion"] + ) + examples = [ + apply_chat_template(example, self.processor) for example in examples + ] prompts = [example["prompt"] for example in examples] completions = [example["completion"] for example in examples] @@ -348,26 +393,51 @@ def torch_call(self, examples: list[dict[str, Any]]) -> dict[str, Any]: ) # Concatenate prompts and completions - prompt_ids, prompt_mask = processed_prompts["input_ids"], processed_prompts["attention_mask"] - completion_ids, completion_mask = processed_completions["input_ids"], processed_completions["attention_mask"] + prompt_ids, prompt_mask = ( + processed_prompts["input_ids"], + processed_prompts["attention_mask"], + ) + completion_ids, completion_mask = ( + processed_completions["input_ids"], + processed_completions["attention_mask"], + ) input_ids = torch.cat((prompt_ids, completion_ids), dim=1) attention_mask = torch.cat((prompt_mask, completion_mask), dim=1) - completion_mask = torch.cat((torch.zeros_like(prompt_mask), completion_mask), dim=1) + completion_mask = torch.cat( + (torch.zeros_like(prompt_mask), completion_mask), dim=1 + ) if "token_type_ids" in processed_prompts: prompt_token_type_ids = processed_prompts["token_type_ids"] completion_token_type_ids = processed_completions["token_type_ids"] - token_type_ids = torch.cat((prompt_token_type_ids, completion_token_type_ids), dim=1) + token_type_ids = torch.cat( + (prompt_token_type_ids, completion_token_type_ids), dim=1 + ) if "mm_token_type_ids" in processed_prompts: prompt_mm_token_type_ids = processed_prompts["mm_token_type_ids"] completion_mm_token_type_ids = processed_completions.get( "mm_token_type_ids", torch.zeros_like(completion_ids) ) - mm_token_type_ids = torch.cat((prompt_mm_token_type_ids, completion_mm_token_type_ids), dim=1) + mm_token_type_ids = torch.cat( + (prompt_mm_token_type_ids, completion_mm_token_type_ids), dim=1 + ) # Flush left to reduce padding - if "token_type_ids" in processed_prompts and "mm_token_type_ids" in processed_prompts: - attention_mask, input_ids, completion_mask, token_type_ids, mm_token_type_ids = flush_left( - attention_mask, input_ids, completion_mask, token_type_ids, mm_token_type_ids + if ( + "token_type_ids" in processed_prompts + and "mm_token_type_ids" in processed_prompts + ): + ( + attention_mask, + input_ids, + completion_mask, + token_type_ids, + mm_token_type_ids, + ) = flush_left( + attention_mask, + input_ids, + completion_mask, + token_type_ids, + mm_token_type_ids, ) elif "token_type_ids" in processed_prompts: attention_mask, input_ids, completion_mask, token_type_ids = flush_left( @@ -378,7 +448,9 @@ def torch_call(self, examples: list[dict[str, Any]]) -> dict[str, Any]: attention_mask, input_ids, completion_mask, mm_token_type_ids ) else: - attention_mask, input_ids, completion_mask = flush_left(attention_mask, input_ids, completion_mask) + attention_mask, input_ids, completion_mask = flush_left( + attention_mask, input_ids, completion_mask + ) # Truncate if necessary if self.max_length is not None: @@ -460,25 +532,54 @@ def add_bos_token_if_needed( rejected_tokens: dict[str, list[int]], ): if bos_token_id is not None: - if prompt_len_input_ids == 0 or bos_token_id != prompt_tokens["prompt_input_ids"][0]: - prompt_tokens["prompt_input_ids"] = [bos_token_id] + prompt_tokens["prompt_input_ids"] - prompt_tokens["prompt_attention_mask"] = [1] + prompt_tokens["prompt_attention_mask"] - if chosen_prompt_len_input_ids == 0 or bos_token_id != chosen_tokens["prompt_input_ids"][0]: - chosen_tokens["prompt_input_ids"] = [bos_token_id] + chosen_tokens["prompt_input_ids"] - chosen_tokens["prompt_attention_mask"] = [1] + chosen_tokens["prompt_attention_mask"] - if rejected_prompt_len_input_ids == 0 or bos_token_id != rejected_tokens["prompt_input_ids"][0]: - rejected_tokens["prompt_input_ids"] = [bos_token_id] + rejected_tokens["prompt_input_ids"] - rejected_tokens["prompt_attention_mask"] = [1] + rejected_tokens["prompt_attention_mask"] + if ( + prompt_len_input_ids == 0 + or bos_token_id != prompt_tokens["prompt_input_ids"][0] + ): + prompt_tokens["prompt_input_ids"] = [bos_token_id] + prompt_tokens[ + "prompt_input_ids" + ] + prompt_tokens["prompt_attention_mask"] = [1] + prompt_tokens[ + "prompt_attention_mask" + ] + if ( + chosen_prompt_len_input_ids == 0 + or bos_token_id != chosen_tokens["prompt_input_ids"][0] + ): + chosen_tokens["prompt_input_ids"] = [bos_token_id] + chosen_tokens[ + "prompt_input_ids" + ] + chosen_tokens["prompt_attention_mask"] = [1] + chosen_tokens[ + "prompt_attention_mask" + ] + if ( + rejected_prompt_len_input_ids == 0 + or bos_token_id != rejected_tokens["prompt_input_ids"][0] + ): + rejected_tokens["prompt_input_ids"] = [bos_token_id] + rejected_tokens[ + "prompt_input_ids" + ] + rejected_tokens["prompt_attention_mask"] = [1] + rejected_tokens[ + "prompt_attention_mask" + ] return prompt_tokens, chosen_tokens, rejected_tokens def add_eos_token_if_needed( - eos_token_id: int, chosen_tokens: dict[str, list[int]], rejected_tokens: dict[str, list[int]] + eos_token_id: int, + chosen_tokens: dict[str, list[int]], + rejected_tokens: dict[str, list[int]], ): - if len(chosen_tokens["input_ids"]) == 0 or eos_token_id != chosen_tokens["input_ids"][-1]: + if ( + len(chosen_tokens["input_ids"]) == 0 + or eos_token_id != chosen_tokens["input_ids"][-1] + ): chosen_tokens["input_ids"].append(eos_token_id) chosen_tokens["attention_mask"].append(1) - if len(rejected_tokens["input_ids"]) == 0 or eos_token_id != rejected_tokens["input_ids"][-1]: + if ( + len(rejected_tokens["input_ids"]) == 0 + or eos_token_id != rejected_tokens["input_ids"][-1] + ): rejected_tokens["input_ids"].append(eos_token_id) rejected_tokens["attention_mask"].append(1) return chosen_tokens, rejected_tokens @@ -503,12 +604,17 @@ def first_true_indices(bools: torch.Tensor, dtype=torch.long) -> torch.Tensor: value is found in a row, returns the length of the row. """ row_len = bools.size(-1) - zero_or_index = row_len * (~bools).type(dtype) + torch.arange(row_len, dtype=dtype, device=bools.device) + zero_or_index = row_len * (~bools).type(dtype) + torch.arange( + row_len, dtype=dtype, device=bools.device + ) return torch.min(zero_or_index, dim=-1).values def get_reward( - model: torch.nn.Module, query_responses: torch.Tensor, pad_token_id: int, context_length: int + model: torch.nn.Module, + query_responses: torch.Tensor, + pad_token_id: int, + context_length: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ Computes the reward logits and the rewards for a given model and query responses. @@ -545,7 +651,11 @@ def get_reward( use_cache=False, # otherwise mistral-based RM would error out ) reward_logits = model.score(output.hidden_states[-1]) - sequence_lengths = first_true_indices(query_responses[:, context_length:] == pad_token_id) - 1 + context_length + sequence_lengths = ( + first_true_indices(query_responses[:, context_length:] == pad_token_id) + - 1 + + context_length + ) # https://github.com/huggingface/transformers/blob/dc68a39c8111217683bf49a4912d0c9018bab33d/src/transformers/models/gpt2/modeling_gpt2.py#L1454 return ( reward_logits, @@ -557,15 +667,19 @@ def get_reward( ) -def prepare_model_for_kbit_training(model, use_gradient_checkpointing=True, gradient_checkpointing_kwargs=None): +def prepare_model_for_kbit_training( + model, use_gradient_checkpointing=True, gradient_checkpointing_kwargs=None +): r""" Prepare a k-bit quantized transformers model for training (PEFT/QLoRA). """ - loaded_in_kbit = getattr(model, "is_loaded_in_8bit", False) or getattr(model, "is_loaded_in_4bit", False) - quant_methods = ["gptq", "aqlm", "eetq", "torchao", "hqq"] - is_quantized = getattr(model, "quantization_method", None) in quant_methods or getattr( - model, "hqq_quantized", False + loaded_in_kbit = getattr(model, "is_loaded_in_8bit", False) or getattr( + model, "is_loaded_in_4bit", False ) + quant_methods = ["gptq", "aqlm", "eetq", "torchao", "hqq"] + is_quantized = getattr( + model, "quantization_method", None + ) in quant_methods or getattr(model, "hqq_quantized", False) if gradient_checkpointing_kwargs is None: gradient_checkpointing_kwargs = {} @@ -588,7 +702,11 @@ def make_inputs_require_grad(module, input, output): supports_gc_kwargs = "gradient_checkpointing_kwargs" in list( inspect.signature(model.gradient_checkpointing_enable).parameters ) - gc_kwargs = {"gradient_checkpointing_kwargs": gradient_checkpointing_kwargs} if supports_gc_kwargs else {} + gc_kwargs = ( + {"gradient_checkpointing_kwargs": gradient_checkpointing_kwargs} + if supports_gc_kwargs + else {} + ) model.gradient_checkpointing_enable(**gc_kwargs) return model @@ -607,7 +725,8 @@ def enable_gradient_checkpointing( gradient_checkpointing_kwargs = gradient_checkpointing_kwargs or {} use_reentrant = ( - "use_reentrant" not in gradient_checkpointing_kwargs or gradient_checkpointing_kwargs["use_reentrant"] + "use_reentrant" not in gradient_checkpointing_kwargs + or gradient_checkpointing_kwargs["use_reentrant"] ) if use_reentrant: @@ -628,7 +747,9 @@ def prepare_peft_model( ) -> PreTrainedModel: """Prepares a model for PEFT training.""" if not is_peft_available(): - raise ImportError("PEFT is required to use a peft model. Run `pip install peft`.") + raise ImportError( + "PEFT is required to use a peft model. Run `pip install peft`." + ) if isinstance(model, PeftModel) and peft_config is not None: raise ValueError( @@ -638,7 +759,9 @@ def prepare_peft_model( ) # Handle quantized models (QLoRA) - is_qlora = getattr(model, "is_loaded_in_4bit", False) or getattr(model, "is_loaded_in_8bit", False) + is_qlora = getattr(model, "is_loaded_in_4bit", False) or getattr( + model, "is_loaded_in_8bit", False + ) is_sharded_qlora = False if getattr(model, "is_loaded_in_4bit", False): @@ -663,7 +786,8 @@ def prepare_peft_model( # Create PEFT model if peft_config is not None: if ( - Version(peft.__version__) >= Version("0.12") # autocast_adapter_dtype introduced in 0.12 + Version(peft.__version__) + >= Version("0.12") # autocast_adapter_dtype introduced in 0.12 and getattr(model, "is_loaded_in_4bit", False) and is_sharded_qlora ): @@ -672,13 +796,19 @@ def prepare_peft_model( model = get_peft_model(model, peft_config) # Handle bf16 casting for 4-bit models - if args.bf16 and getattr(model, "is_loaded_in_4bit", False) and not is_sharded_qlora: + if ( + args.bf16 + and getattr(model, "is_loaded_in_4bit", False) + and not is_sharded_qlora + ): peft_module_casting_to_bf16(model) return model -def pad_to_length(tensor: torch.Tensor, length: int, pad_value: int | float, dim: int = -1) -> torch.Tensor: +def pad_to_length( + tensor: torch.Tensor, length: int, pad_value: int | float, dim: int = -1 +) -> torch.Tensor: if tensor.size(dim) >= length: return tensor else: @@ -687,7 +817,8 @@ def pad_to_length(tensor: torch.Tensor, length: int, pad_value: int | float, dim return torch.cat( [ tensor, - pad_value * torch.ones(*pad_size, dtype=tensor.dtype, device=tensor.device), + pad_value + * torch.ones(*pad_size, dtype=tensor.dtype, device=tensor.device), ], dim=dim, ) @@ -799,7 +930,9 @@ def create_reference_model( param.requires_grad = False if pattern is not None and len(unshared_param_list) == 0: - logging.warning("Pattern passed or found, but no layers matched in the model. Check for a typo.") + logging.warning( + "Pattern passed or found, but no layers matched in the model. Check for a typo." + ) return ref_model.eval() @@ -844,7 +977,9 @@ def truncate_dataset( def truncate(examples): truncated_columns = [] for column in examples.columns: - if pyarrow.types.is_list(column.type) or pyarrow.types.is_large_list(column.type): + if pyarrow.types.is_list(column.type) or pyarrow.types.is_large_list( + column.type + ): column = pc.list_slice(column, 0, max_length) truncated_columns.append(column) return pa.Table.from_arrays(truncated_columns, names=examples.column_names) From fd3be85352a49dc578c5861b20aaf3ed880f6e55 Mon Sep 17 00:00:00 2001 From: Strongich Date: Wed, 8 Apr 2026 18:29:59 +0200 Subject: [PATCH 11/39] precommit ruff format --- trl/experimental/utils.py | 214 +++++++++----------------------------- 1 file changed, 52 insertions(+), 162 deletions(-) diff --git a/trl/experimental/utils.py b/trl/experimental/utils.py index 5526439717a..3c18412beb9 100644 --- a/trl/experimental/utils.py +++ b/trl/experimental/utils.py @@ -73,9 +73,7 @@ def __call__(self, features: list[dict[str, Any]]) -> dict[str, Any]: # first, pad everything to the same length padded_batch = {} for k in features[0].keys(): - if k.endswith( - ("_input_ids", "_attention_mask", "_labels", "_pixel_values") - ): + if k.endswith(("_input_ids", "_attention_mask", "_labels", "_pixel_values")): if self.is_encoder_decoder: to_pad = [torch.LongTensor(ex[k]) for ex in features] @@ -89,15 +87,11 @@ def __call__(self, features: list[dict[str, Any]]) -> dict[str, Any]: padding_value = self.pad_token_id elif k.endswith("_attention_mask"): padding_value = 0 - elif k.startswith(("chosen", "rejected", "completion")) or ( - "decoder" in k - ): + elif k.startswith(("chosen", "rejected", "completion")) or ("decoder" in k): padding_value = -100 else: raise ValueError(f"Unexpected key in batch '{k}'") - padded_batch[k] = pad_sequence( - to_pad, batch_first=True, padding_value=padding_value - ) + padded_batch[k] = pad_sequence(to_pad, batch_first=True, padding_value=padding_value) else: # Set padding value based on the key if k.endswith("_input_ids"): @@ -125,17 +119,13 @@ def __call__(self, features: list[dict[str, Any]]) -> dict[str, Any]: # Set the dtype if k.endswith("_pixel_values"): - dtype = ( - torch.float32 - ) # will be downcasted if necessary by the Trainer + dtype = torch.float32 # will be downcasted if necessary by the Trainer else: dtype = torch.int64 # Convert to tensor and pad to_pad = [torch.tensor(ex[k], dtype=dtype) for ex in features] - padded_batch[k] = pad( - to_pad, padding_value=padding_value, padding_side=padding_side - ) + padded_batch[k] = pad(to_pad, padding_value=padding_value, padding_side=padding_side) elif k.endswith("_logps"): # the cached reference model logprobs padded_batch[k] = torch.tensor([ex[k] for ex in features]) @@ -159,9 +149,7 @@ class DataCollatorForChatML: def __post_init__(self): if self.tokenizer.pad_token_id is None: - raise ValueError( - "The tokenizer does not have a pad token. Please set `pad_token_id` in the tokenizer." - ) + raise ValueError("The tokenizer does not have a pad token. Please set `pad_token_id` in the tokenizer.") if self.max_length is None: # set a sensible default self.max_length = min(self.tokenizer.model_max_length, 1024) @@ -201,11 +189,7 @@ def __call__(self, examples: list[dict[str, Any]]) -> dict[str, torch.Tensor]: if offsets is not None: prompt_char_len = len(formatted_prompt) completion_start_idx_full = next( - ( - idx - for idx, (start, _) in enumerate(offsets) - if start >= prompt_char_len - ), + (idx for idx, (start, _) in enumerate(offsets) if start >= prompt_char_len), len(message_input_ids_full), ) else: @@ -219,25 +203,16 @@ def __call__(self, examples: list[dict[str, Any]]) -> dict[str, torch.Tensor]: completion_start_idx_full = len(tokenized_prompt_full["input_ids"]) prompt_tokens_full = message_input_ids_full[:completion_start_idx_full] - completion_input_ids_full = message_input_ids_full[ - completion_start_idx_full: - ] - - if ( - self.max_length is not None - and len(message_input_ids_full) > self.max_length - ): + completion_input_ids_full = message_input_ids_full[completion_start_idx_full:] + + if self.max_length is not None and len(message_input_ids_full) > self.max_length: completion_ids = completion_input_ids_full if len(completion_ids) >= self.max_length: completion_ids = completion_ids[-self.max_length :] prompt_ids = [] else: max_prompt_tokens = self.max_length - len(completion_ids) - prompt_ids = ( - prompt_tokens_full[-max_prompt_tokens:] - if max_prompt_tokens > 0 - else [] - ) + prompt_ids = prompt_tokens_full[-max_prompt_tokens:] if max_prompt_tokens > 0 else [] message_input_ids = prompt_ids + completion_ids else: message_input_ids = message_input_ids_full @@ -274,30 +249,20 @@ def __call__(self, examples: list[dict[str, Any]]) -> dict[str, torch.Tensor]: # convert to list of tensors and pad input_ids = [torch.tensor(ids, dtype=torch.long) for ids in input_ids] - attention_mask = [ - torch.tensor(mask, dtype=torch.long) for mask in attention_mask - ] + attention_mask = [torch.tensor(mask, dtype=torch.long) for mask in attention_mask] labels = [torch.tensor(label, dtype=torch.long) for label in labels] - input_ids = pad( - input_ids, padding_side="left", padding_value=self.tokenizer.pad_token_id - ) + input_ids = pad(input_ids, padding_side="left", padding_value=self.tokenizer.pad_token_id) attention_mask = pad(attention_mask, padding_side="left", padding_value=0) labels = pad(labels, padding_side="left", padding_value=self.ignore_index) - prompts_input_ids = [ - torch.tensor(ids, dtype=torch.long) for ids in prompts_input_ids - ] - prompt_attention_mask = [ - torch.tensor(mask, dtype=torch.long) for mask in prompt_attention_mask - ] + prompts_input_ids = [torch.tensor(ids, dtype=torch.long) for ids in prompts_input_ids] + prompt_attention_mask = [torch.tensor(mask, dtype=torch.long) for mask in prompt_attention_mask] prompts_input_ids = pad( prompts_input_ids, padding_side="left", padding_value=self.tokenizer.pad_token_id, ) - prompt_attention_mask = pad( - prompt_attention_mask, padding_side="left", padding_value=0 - ) + prompt_attention_mask = pad(prompt_attention_mask, padding_side="left", padding_value=0) return { "input_ids": input_ids, @@ -362,15 +327,9 @@ def torch_call(self, examples: list[dict[str, Any]]) -> dict[str, Any]: # Apply chat template for conversational data if is_conversational(examples[0]): for example in examples: - example["prompt"] = prepare_multimodal_messages( - example["prompt"], images=example["images"] - ) - example["completion"] = prepare_multimodal_messages( - example["completion"] - ) - examples = [ - apply_chat_template(example, self.processor) for example in examples - ] + example["prompt"] = prepare_multimodal_messages(example["prompt"], images=example["images"]) + example["completion"] = prepare_multimodal_messages(example["completion"]) + examples = [apply_chat_template(example, self.processor) for example in examples] prompts = [example["prompt"] for example in examples] completions = [example["completion"] for example in examples] @@ -403,29 +362,20 @@ def torch_call(self, examples: list[dict[str, Any]]) -> dict[str, Any]: ) input_ids = torch.cat((prompt_ids, completion_ids), dim=1) attention_mask = torch.cat((prompt_mask, completion_mask), dim=1) - completion_mask = torch.cat( - (torch.zeros_like(prompt_mask), completion_mask), dim=1 - ) + completion_mask = torch.cat((torch.zeros_like(prompt_mask), completion_mask), dim=1) if "token_type_ids" in processed_prompts: prompt_token_type_ids = processed_prompts["token_type_ids"] completion_token_type_ids = processed_completions["token_type_ids"] - token_type_ids = torch.cat( - (prompt_token_type_ids, completion_token_type_ids), dim=1 - ) + token_type_ids = torch.cat((prompt_token_type_ids, completion_token_type_ids), dim=1) if "mm_token_type_ids" in processed_prompts: prompt_mm_token_type_ids = processed_prompts["mm_token_type_ids"] completion_mm_token_type_ids = processed_completions.get( "mm_token_type_ids", torch.zeros_like(completion_ids) ) - mm_token_type_ids = torch.cat( - (prompt_mm_token_type_ids, completion_mm_token_type_ids), dim=1 - ) + mm_token_type_ids = torch.cat((prompt_mm_token_type_ids, completion_mm_token_type_ids), dim=1) # Flush left to reduce padding - if ( - "token_type_ids" in processed_prompts - and "mm_token_type_ids" in processed_prompts - ): + if "token_type_ids" in processed_prompts and "mm_token_type_ids" in processed_prompts: ( attention_mask, input_ids, @@ -448,9 +398,7 @@ def torch_call(self, examples: list[dict[str, Any]]) -> dict[str, Any]: attention_mask, input_ids, completion_mask, mm_token_type_ids ) else: - attention_mask, input_ids, completion_mask = flush_left( - attention_mask, input_ids, completion_mask - ) + attention_mask, input_ids, completion_mask = flush_left(attention_mask, input_ids, completion_mask) # Truncate if necessary if self.max_length is not None: @@ -532,36 +480,15 @@ def add_bos_token_if_needed( rejected_tokens: dict[str, list[int]], ): if bos_token_id is not None: - if ( - prompt_len_input_ids == 0 - or bos_token_id != prompt_tokens["prompt_input_ids"][0] - ): - prompt_tokens["prompt_input_ids"] = [bos_token_id] + prompt_tokens[ - "prompt_input_ids" - ] - prompt_tokens["prompt_attention_mask"] = [1] + prompt_tokens[ - "prompt_attention_mask" - ] - if ( - chosen_prompt_len_input_ids == 0 - or bos_token_id != chosen_tokens["prompt_input_ids"][0] - ): - chosen_tokens["prompt_input_ids"] = [bos_token_id] + chosen_tokens[ - "prompt_input_ids" - ] - chosen_tokens["prompt_attention_mask"] = [1] + chosen_tokens[ - "prompt_attention_mask" - ] - if ( - rejected_prompt_len_input_ids == 0 - or bos_token_id != rejected_tokens["prompt_input_ids"][0] - ): - rejected_tokens["prompt_input_ids"] = [bos_token_id] + rejected_tokens[ - "prompt_input_ids" - ] - rejected_tokens["prompt_attention_mask"] = [1] + rejected_tokens[ - "prompt_attention_mask" - ] + if prompt_len_input_ids == 0 or bos_token_id != prompt_tokens["prompt_input_ids"][0]: + prompt_tokens["prompt_input_ids"] = [bos_token_id] + prompt_tokens["prompt_input_ids"] + prompt_tokens["prompt_attention_mask"] = [1] + prompt_tokens["prompt_attention_mask"] + if chosen_prompt_len_input_ids == 0 or bos_token_id != chosen_tokens["prompt_input_ids"][0]: + chosen_tokens["prompt_input_ids"] = [bos_token_id] + chosen_tokens["prompt_input_ids"] + chosen_tokens["prompt_attention_mask"] = [1] + chosen_tokens["prompt_attention_mask"] + if rejected_prompt_len_input_ids == 0 or bos_token_id != rejected_tokens["prompt_input_ids"][0]: + rejected_tokens["prompt_input_ids"] = [bos_token_id] + rejected_tokens["prompt_input_ids"] + rejected_tokens["prompt_attention_mask"] = [1] + rejected_tokens["prompt_attention_mask"] return prompt_tokens, chosen_tokens, rejected_tokens @@ -570,16 +497,10 @@ def add_eos_token_if_needed( chosen_tokens: dict[str, list[int]], rejected_tokens: dict[str, list[int]], ): - if ( - len(chosen_tokens["input_ids"]) == 0 - or eos_token_id != chosen_tokens["input_ids"][-1] - ): + if len(chosen_tokens["input_ids"]) == 0 or eos_token_id != chosen_tokens["input_ids"][-1]: chosen_tokens["input_ids"].append(eos_token_id) chosen_tokens["attention_mask"].append(1) - if ( - len(rejected_tokens["input_ids"]) == 0 - or eos_token_id != rejected_tokens["input_ids"][-1] - ): + if len(rejected_tokens["input_ids"]) == 0 or eos_token_id != rejected_tokens["input_ids"][-1]: rejected_tokens["input_ids"].append(eos_token_id) rejected_tokens["attention_mask"].append(1) return chosen_tokens, rejected_tokens @@ -604,9 +525,7 @@ def first_true_indices(bools: torch.Tensor, dtype=torch.long) -> torch.Tensor: value is found in a row, returns the length of the row. """ row_len = bools.size(-1) - zero_or_index = row_len * (~bools).type(dtype) + torch.arange( - row_len, dtype=dtype, device=bools.device - ) + zero_or_index = row_len * (~bools).type(dtype) + torch.arange(row_len, dtype=dtype, device=bools.device) return torch.min(zero_or_index, dim=-1).values @@ -651,11 +570,7 @@ def get_reward( use_cache=False, # otherwise mistral-based RM would error out ) reward_logits = model.score(output.hidden_states[-1]) - sequence_lengths = ( - first_true_indices(query_responses[:, context_length:] == pad_token_id) - - 1 - + context_length - ) + sequence_lengths = first_true_indices(query_responses[:, context_length:] == pad_token_id) - 1 + context_length # https://github.com/huggingface/transformers/blob/dc68a39c8111217683bf49a4912d0c9018bab33d/src/transformers/models/gpt2/modeling_gpt2.py#L1454 return ( reward_logits, @@ -667,19 +582,15 @@ def get_reward( ) -def prepare_model_for_kbit_training( - model, use_gradient_checkpointing=True, gradient_checkpointing_kwargs=None -): +def prepare_model_for_kbit_training(model, use_gradient_checkpointing=True, gradient_checkpointing_kwargs=None): r""" Prepare a k-bit quantized transformers model for training (PEFT/QLoRA). """ - loaded_in_kbit = getattr(model, "is_loaded_in_8bit", False) or getattr( - model, "is_loaded_in_4bit", False - ) + loaded_in_kbit = getattr(model, "is_loaded_in_8bit", False) or getattr(model, "is_loaded_in_4bit", False) quant_methods = ["gptq", "aqlm", "eetq", "torchao", "hqq"] - is_quantized = getattr( - model, "quantization_method", None - ) in quant_methods or getattr(model, "hqq_quantized", False) + is_quantized = getattr(model, "quantization_method", None) in quant_methods or getattr( + model, "hqq_quantized", False + ) if gradient_checkpointing_kwargs is None: gradient_checkpointing_kwargs = {} @@ -702,11 +613,7 @@ def make_inputs_require_grad(module, input, output): supports_gc_kwargs = "gradient_checkpointing_kwargs" in list( inspect.signature(model.gradient_checkpointing_enable).parameters ) - gc_kwargs = ( - {"gradient_checkpointing_kwargs": gradient_checkpointing_kwargs} - if supports_gc_kwargs - else {} - ) + gc_kwargs = {"gradient_checkpointing_kwargs": gradient_checkpointing_kwargs} if supports_gc_kwargs else {} model.gradient_checkpointing_enable(**gc_kwargs) return model @@ -725,8 +632,7 @@ def enable_gradient_checkpointing( gradient_checkpointing_kwargs = gradient_checkpointing_kwargs or {} use_reentrant = ( - "use_reentrant" not in gradient_checkpointing_kwargs - or gradient_checkpointing_kwargs["use_reentrant"] + "use_reentrant" not in gradient_checkpointing_kwargs or gradient_checkpointing_kwargs["use_reentrant"] ) if use_reentrant: @@ -747,9 +653,7 @@ def prepare_peft_model( ) -> PreTrainedModel: """Prepares a model for PEFT training.""" if not is_peft_available(): - raise ImportError( - "PEFT is required to use a peft model. Run `pip install peft`." - ) + raise ImportError("PEFT is required to use a peft model. Run `pip install peft`.") if isinstance(model, PeftModel) and peft_config is not None: raise ValueError( @@ -759,9 +663,7 @@ def prepare_peft_model( ) # Handle quantized models (QLoRA) - is_qlora = getattr(model, "is_loaded_in_4bit", False) or getattr( - model, "is_loaded_in_8bit", False - ) + is_qlora = getattr(model, "is_loaded_in_4bit", False) or getattr(model, "is_loaded_in_8bit", False) is_sharded_qlora = False if getattr(model, "is_loaded_in_4bit", False): @@ -786,8 +688,7 @@ def prepare_peft_model( # Create PEFT model if peft_config is not None: if ( - Version(peft.__version__) - >= Version("0.12") # autocast_adapter_dtype introduced in 0.12 + Version(peft.__version__) >= Version("0.12") # autocast_adapter_dtype introduced in 0.12 and getattr(model, "is_loaded_in_4bit", False) and is_sharded_qlora ): @@ -796,19 +697,13 @@ def prepare_peft_model( model = get_peft_model(model, peft_config) # Handle bf16 casting for 4-bit models - if ( - args.bf16 - and getattr(model, "is_loaded_in_4bit", False) - and not is_sharded_qlora - ): + if args.bf16 and getattr(model, "is_loaded_in_4bit", False) and not is_sharded_qlora: peft_module_casting_to_bf16(model) return model -def pad_to_length( - tensor: torch.Tensor, length: int, pad_value: int | float, dim: int = -1 -) -> torch.Tensor: +def pad_to_length(tensor: torch.Tensor, length: int, pad_value: int | float, dim: int = -1) -> torch.Tensor: if tensor.size(dim) >= length: return tensor else: @@ -817,8 +712,7 @@ def pad_to_length( return torch.cat( [ tensor, - pad_value - * torch.ones(*pad_size, dtype=tensor.dtype, device=tensor.device), + pad_value * torch.ones(*pad_size, dtype=tensor.dtype, device=tensor.device), ], dim=dim, ) @@ -930,9 +824,7 @@ def create_reference_model( param.requires_grad = False if pattern is not None and len(unshared_param_list) == 0: - logging.warning( - "Pattern passed or found, but no layers matched in the model. Check for a typo." - ) + logging.warning("Pattern passed or found, but no layers matched in the model. Check for a typo.") return ref_model.eval() @@ -977,9 +869,7 @@ def truncate_dataset( def truncate(examples): truncated_columns = [] for column in examples.columns: - if pyarrow.types.is_list(column.type) or pyarrow.types.is_large_list( - column.type - ): + if pyarrow.types.is_list(column.type) or pyarrow.types.is_large_list(column.type): column = pc.list_slice(column, 0, max_length) truncated_columns.append(column) return pa.Table.from_arrays(truncated_columns, names=examples.column_names) From 719c560c985ee6e77629b9a0802793fd5a818e0f Mon Sep 17 00:00:00 2001 From: Strongich Date: Tue, 21 Apr 2026 16:41:06 +0200 Subject: [PATCH 12/39] fix: pre-duplicate prompts for num_generations > 1 in VLM vLLM path --- tests/experimental/test_gold_trainer.py | 98 +++++++++++++++++++++++++ trl/experimental/gold/gold_trainer.py | 23 ++++-- 2 files changed, 114 insertions(+), 7 deletions(-) diff --git a/tests/experimental/test_gold_trainer.py b/tests/experimental/test_gold_trainer.py index 76df35cc03c..0789db066ea 100644 --- a/tests/experimental/test_gold_trainer.py +++ b/tests/experimental/test_gold_trainer.py @@ -1743,3 +1743,101 @@ def fake_sft_init( # Identity collator and VLM collator should still be set assert trainer.data_collator is identity assert trainer._vlm_collator is not None + + +def test_on_policy_vlm_vllm_duplicates_prompts_for_num_generations(monkeypatch): + """Regression: in the vLLM path, `_generate_on_policy_vlm_raw` must pre-duplicate prompts + `num_generations` times to satisfy `vllm_generation.generate`'s contract (which returns one completion per input + prompt entry; colocate mode hardcodes `n=1`). Without duplication, the redistribution loop raised `IndexError` for + `num_generations > 1`. + """ + num_generations = 3 + num_prompts_per_slice = 2 + num_slices = 2 + + trainer = GOLDTrainer.__new__(GOLDTrainer) + trainer.accelerator = SimpleNamespace(device=torch.device("cpu")) + trainer.use_vllm = True + trainer.num_generations = num_generations + trainer.state = SimpleNamespace(global_step=0) + trainer._last_vllm_sync_step = -1 + trainer.vllm_sync_frequency = 1 + trainer.generation_config = SimpleNamespace(max_new_tokens=16) + trainer._buffered_inputs = {} + trainer._buffered_text_logs = {} + trainer._teacher_processor = None + + class StubProcessor: + @staticmethod + def apply_chat_template(conversation, add_generation_prompt, tokenize, return_dict, padding): + return { + "input_ids": [[1, 2, 3] for _ in conversation], + "attention_mask": [[1, 1, 1] for _ in conversation], + } + + @staticmethod + def batch_decode(ids, skip_special_tokens): + return [f"prompt_{i}" for i in range(len(ids))] + + @staticmethod + def decode(ids, skip_special_tokens, clean_up_tokenization_spaces): + return f"comp_{ids[0]}" + + trainer.processing_class = StubProcessor + + received = {} + + class StubVLLMGeneration: + def sync_weights(self): + pass + + def generate(self, prompts, images, num_generations): + received["n_prompts"] = len(prompts) + received["n_images"] = len(images) if images is not None else None + completion_ids = [[100 + i] for i in range(len(prompts))] + return None, completion_ids, None, None + + trainer.vllm_generation = StubVLLMGeneration() + + collated_per_call = [] + + def stub_collator(synthetic_examples): + collated_per_call.append(list(synthetic_examples)) + return {"input_ids": torch.zeros(len(synthetic_examples), 1, dtype=torch.long)} + + trainer._vlm_collator = stub_collator + + class FakeImage: + def __init__(self, tag): + self.tag = tag + + raw_slices = [ + [ + {"prompt": [{"role": "user", "content": f"q{slice_idx}_{i}"}], "image": FakeImage(f"{slice_idx}_{i}")} + for i in range(num_prompts_per_slice) + ] + for slice_idx in range(num_slices) + ] + on_policy_indices = list(range(num_slices)) + + # Bypass multimodal-message helper; its exact shape is irrelevant to this regression. + monkeypatch.setattr(gold_trainer_module, "prepare_multimodal_messages", lambda prompt, images: prompt) + + trainer._generate_on_policy_vlm_raw(raw_slices, on_policy_indices) + + total_unique_prompts = num_slices * num_prompts_per_slice + # The fix: prompts (and images) are pre-duplicated `num_generations` times before + # being handed to vllm_generation.generate. + assert received["n_prompts"] == total_unique_prompts * num_generations + assert received["n_images"] == total_unique_prompts * num_generations + + # Each slice ends up with `num_prompts_per_slice * num_generations` synthetic examples. + assert len(collated_per_call) == num_slices + for synthetic in collated_per_call: + assert len(synthetic) == num_prompts_per_slice * num_generations + + # Buffers populated for every on-policy slice without IndexError. + for slice_idx in on_policy_indices: + assert slice_idx in trainer._buffered_inputs + _, completion_texts = trainer._buffered_text_logs[slice_idx] + assert len(completion_texts) == num_prompts_per_slice * num_generations diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index ac2e574409c..a7eeda088b4 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -1516,10 +1516,17 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in self.vllm_generation.sync_weights() self._last_vllm_sync_step = self.state.global_step - # Pass None for images if all entries are None - generate_images = all_images if any(img is not None for img in all_images) else None + # `vllm_generation.generate` returns one completion per input prompt entry and expects the + # caller to pre-duplicate prompts `num_generations` times (same contract as GRPOTrainer). + # Without this duplication, colocate mode (which hardcodes n=1) yields only `len(prompts)` + # completions while the redistribution below expects `len(prompts) * num_generations`. + dup_prompt_ids = [ids for ids in all_prompt_ids for _ in range(self.num_generations)] + if any(img is not None for img in all_images): + generate_images = [img for img in all_images for _ in range(self.num_generations)] + else: + generate_images = None _, completion_ids, _, _ = self.vllm_generation.generate( - prompts=all_prompt_ids, + prompts=dup_prompt_ids, images=generate_images, num_generations=self.num_generations, ) @@ -1534,8 +1541,9 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in self.processing_class.decode(comp_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False) ) - # Redistribute completions to slices. With num_generations > 1, each prompt produces - # multiple completions, so we repeat each raw example/prompt/image to match. + # Redistribute completions to slices. Completions now align 1:1 with the duplicated inputs, + # so `comp_idx` is a single running counter; raw example/prompt/image entries still reference + # the original unique example `i` since those are shared across the `num_generations` copies. slice_completions = {idx: [] for idx in on_policy_indices} slice_raw = {idx: [] for idx in on_policy_indices} slice_images = {idx: [] for idx in on_policy_indices} @@ -1543,15 +1551,16 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in slice_prompts_text = {idx: [] for idx in on_policy_indices} slice_prompts_text_special = {idx: [] for idx in on_policy_indices} + comp_idx = 0 for i, slice_idx in enumerate(local_slice_indices): - for g in range(self.num_generations): - comp_idx = i * self.num_generations + g + for _ in range(self.num_generations): slice_completions[slice_idx].append(all_completion_texts[comp_idx]) slice_raw[slice_idx].append(all_raw_examples[i]) slice_images[slice_idx].append(all_images[i]) slice_prompts[slice_idx].append(all_prompts[i]) slice_prompts_text[slice_idx].append(all_prompts_text[i]) slice_prompts_text_special[slice_idx].append(all_prompts_text_with_special[i]) + comp_idx += 1 for slice_idx in on_policy_indices: completion_texts = slice_completions[slice_idx] From d3dcf08fdc936827df1cd92c8ee77e93e6b648ab Mon Sep 17 00:00:00 2001 From: Strongich Date: Tue, 21 Apr 2026 16:54:06 +0200 Subject: [PATCH 13/39] wrap on-policy VLM completions as content blocks --- trl/experimental/gold/gold_trainer.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index a7eeda088b4..78bb487abd8 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -1572,7 +1572,11 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in synthetic_examples = [] for i, example in enumerate(raw_for_slice): synthetic = dict(example) - synthetic["completion"] = [{"role": "assistant", "content": completion_texts[i]}] + # Wrap as content blocks so VLM chat templates (e.g. SmolVLM) that index + # `message.content[0]` can render the synthetic assistant turn. + synthetic["completion"] = [ + {"role": "assistant", "content": [{"type": "text", "text": completion_texts[i]}]} + ] synthetic_examples.append(synthetic) # Collate synthetic examples to get pixel_values + properly tokenized input_ids/labels From 23a294f47b2ddcd0950266eab5d07de850b33703 Mon Sep 17 00:00:00 2001 From: Strongich Date: Tue, 21 Apr 2026 17:25:56 +0200 Subject: [PATCH 14/39] Strip student chat-template markers from ULD text fields in VLM collator --- tests/experimental/test_gold_trainer.py | 41 +++++++++++++++++++++++++ trl/experimental/utils.py | 29 +++++++++++++++-- 2 files changed, 67 insertions(+), 3 deletions(-) diff --git a/tests/experimental/test_gold_trainer.py b/tests/experimental/test_gold_trainer.py index 0789db066ea..ff3edee07fe 100644 --- a/tests/experimental/test_gold_trainer.py +++ b/tests/experimental/test_gold_trainer.py @@ -1248,6 +1248,47 @@ def test_vlm_collator_label_masking(smolvlm_processor, vlm_examples): assert prompt_positions.any(), "Each example must have masked prompt tokens" +def test_vlm_collator_original_text_is_untemplated(smolvlm_processor, vlm_examples): + """`original_*_text` must be free of the student's chat-template markers. + + Cross-tokenizer ULD distillation re-renders the prompt through the teacher's chat template and concatenates the + stored completion. If the stored completion still carries the student's special tokens (e.g. ``<|im_end|>``, role + headers), the teacher tokenizer will tokenize them as regular text, producing spurious teacher tokens and incorrect + teacher logits. + """ + collator = DataCollatorForVisionLanguageChatML(processor=smolvlm_processor, max_length=2048) + # Take a deep copy because the collator mutates examples in-place. + import copy + + batch = collator(copy.deepcopy(vlm_examples)) + + student_tokenizer = smolvlm_processor.tokenizer + student_specials = [tok for tok in student_tokenizer.all_special_tokens if tok and tok.strip()] + + assert "original_prompt_text" in batch + assert "original_completion_text" in batch + assert len(batch["original_prompt_text"]) == len(vlm_examples) + assert len(batch["original_completion_text"]) == len(vlm_examples) + + expected_assistant_texts = _get_assistant_texts(vlm_examples) + for raw_completion, assistant_text in zip( + batch["original_completion_text"], expected_assistant_texts, strict=True + ): + # The raw text must still contain the underlying assistant content... + assert assistant_text.strip() in raw_completion + # ...but must not include any of the student's chat-template special tokens. + for special in student_specials: + assert special not in raw_completion, ( + f"original_completion_text leaked student special token {special!r}: {raw_completion!r}" + ) + + for raw_prompt in batch["original_prompt_text"]: + for special in student_specials: + assert special not in raw_prompt, ( + f"original_prompt_text leaked student special token {special!r}: {raw_prompt!r}" + ) + + def test_gold_trainer_init_rejects_non_vlm_teacher(monkeypatch): """GOLDTrainer should raise ValueError when the student is a VLM but the teacher is not.""" diff --git a/trl/experimental/utils.py b/trl/experimental/utils.py index 3c18412beb9..e1e667b7fe9 100644 --- a/trl/experimental/utils.py +++ b/trl/experimental/utils.py @@ -324,6 +324,26 @@ def torch_call(self, examples: list[dict[str, Any]]) -> dict[str, Any]: if all(img_list == [] for img_list in images): images = None + # Capture raw prompt/completion text before `apply_chat_template` mutates the examples. + def _raw_text_from_messages(messages_or_str: Any) -> str: + if isinstance(messages_or_str, str): + return messages_or_str + parts: list[str] = [] + for turn in messages_or_str: + content = turn.get("content", "") + if isinstance(content, str): + parts.append(content) + continue + for block in content: + if isinstance(block, dict) and block.get("type") == "text": + parts.append(block.get("text", "")) + elif isinstance(block, str): + parts.append(block) + return "".join(parts) + + raw_prompt_texts = [_raw_text_from_messages(example["prompt"]) for example in examples] + raw_completion_texts = [_raw_text_from_messages(example["completion"]) for example in examples] + # Apply chat template for conversational data if is_conversational(examples[0]): for example in examples: @@ -431,9 +451,12 @@ def torch_call(self, examples: list[dict[str, Any]]) -> dict[str, Any]: output["prompts"] = prompt_ids output["prompt_attention_mask"] = prompt_mask - # GOLD-specific: raw text for ULD cross-tokenizer distillation - output["original_prompt_text"] = prompts - output["original_completion_text"] = completions + # GOLD-specific: raw text for ULD cross-tokenizer distillation. + # These must be the untemplated text (no student chat-template markers) so the + # teacher can re-render the prompt through its own chat template and tokenize the + # completion cleanly. + output["original_prompt_text"] = raw_prompt_texts + output["original_completion_text"] = raw_completion_texts return output From 9c661c6fcd2dd0435f9ce406cf4cc24e56bd2041 Mon Sep 17 00:00:00 2001 From: Strongich Date: Tue, 21 Apr 2026 18:46:27 +0200 Subject: [PATCH 15/39] ignore empty VLM label rows from zeroing JSD loss --- tests/experimental/test_gold_trainer.py | 32 +++++++++++++++++++++++++ trl/experimental/gold/gold_trainer.py | 15 ++++++++++-- 2 files changed, 45 insertions(+), 2 deletions(-) diff --git a/tests/experimental/test_gold_trainer.py b/tests/experimental/test_gold_trainer.py index ff3edee07fe..8d6bbd4008d 100644 --- a/tests/experimental/test_gold_trainer.py +++ b/tests/experimental/test_gold_trainer.py @@ -824,6 +824,38 @@ def test_generalized_jsd_loss_accepts_probability_inputs(): torch.testing.assert_close(loss, expected) +def test_generalized_jsd_loss_returns_zero_when_all_labels_are_ignored(): + student_logits = torch.tensor([[[0.5, 1.5], [1.0, 0.0]]]) + teacher_logits = torch.tensor([[[1.0, 0.0], [0.5, 1.5]]]) + labels = torch.full((1, 2), -100) + + loss = GOLDTrainer.generalized_jsd_loss(student_logits, teacher_logits, labels=labels) + + torch.testing.assert_close(loss, torch.tensor(0.0)) + + +def test_vlm_prompt_length_ignores_rows_without_completion_labels(): + labels = torch.tensor( + [ + [-100, -100, -100, -100, -100], + [-100, -100, 10, 11, -100], + [-100, -100, -100, 12, 13], + ] + ) + + prompt_length = GOLDTrainer._get_min_completion_start_from_labels(labels) + + assert prompt_length == 2 + + +def test_vlm_prompt_length_uses_sequence_length_when_all_labels_are_ignored(): + labels = torch.full((2, 4), -100) + + prompt_length = GOLDTrainer._get_min_completion_start_from_labels(labels) + + assert prompt_length == labels.shape[1] + + def test_uldloss_handles_llama_student_qwen_teacher_sequence(llama_tokenizer, qwen_tokenizer): config = build_config( uld_use_hybrid_loss=True, diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index 78bb487abd8..ddea1899ddd 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -1229,6 +1229,15 @@ def _build_sequence_batch( return new_attention_mask, new_labels + @staticmethod + def _get_min_completion_start_from_labels(labels: torch.Tensor) -> int: + """Return the earliest completion start, ignoring rows with no completion labels.""" + completion_mask = labels != -100 + rows_with_completion = completion_mask.any(dim=1) + if not rows_with_completion.any(): + return labels.shape[1] + return completion_mask[rows_with_completion].long().argmax(dim=1).min().item() + @profiling_decorator def _fill_buffer(self, generation_batch: dict[str, torch.Tensor | Any] | list[dict], buffer_steps: int): if self._vlm_collator is not None: @@ -2040,6 +2049,8 @@ def generalized_jsd_loss( # Masking if labels is not None: mask = labels != -100 + if not mask.any(): + return jsd.sum() * 0.0 jsd = jsd[mask] # Apply reduction @@ -2165,7 +2176,7 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N # For VLMs, prompts are left-padded but input_ids are flushed left, so prompts.shape[1] # would overcount. Derive prompt length from the flushed labels instead. if self._is_vlm: - student_prompt_length = (inputs["labels"] != -100).long().argmax(dim=1).min().item() + student_prompt_length = self._get_min_completion_start_from_labels(inputs["labels"]) else: student_prompt_length = inputs["prompts"].shape[1] shifted_student_logits = outputs_student.logits[:, student_prompt_length - 1 : -1, :] @@ -2250,7 +2261,7 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N # Using the same prompt_lengths for teacher and student, since JSD can only be # used with same-family VLMs (shared tokenizer). if self._is_vlm: - prompt_lengths = (inputs["labels"] != -100).long().argmax(dim=1).min().item() + prompt_lengths = self._get_min_completion_start_from_labels(inputs["labels"]) else: prompt_lengths = inputs["prompts"].shape[1] shifted_student_logits = outputs_student.logits[:, prompt_lengths - 1 : -1, :] From 5c44e977bc8d28bd276a341ceeeb9e88a1081055 Mon Sep 17 00:00:00 2001 From: Strongich Date: Sat, 16 May 2026 18:30:38 +0200 Subject: [PATCH 16/39] switch to torch_dtype --- examples/scripts/gold_vlm.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/examples/scripts/gold_vlm.py b/examples/scripts/gold_vlm.py index f50ad67a691..e2d87e6cd07 100644 --- a/examples/scripts/gold_vlm.py +++ b/examples/scripts/gold_vlm.py @@ -104,8 +104,12 @@ def convert_to_rgb(example): # ────────────────────────────────────────────── # Models # ────────────────────────────────────────────── - student_model = AutoModelForImageTextToText.from_pretrained(cli_args.student_model_name, dtype=torch.bfloat16) - teacher_model = AutoModelForImageTextToText.from_pretrained(cli_args.teacher_model_name, dtype=torch.bfloat16) + student_model = AutoModelForImageTextToText.from_pretrained( + cli_args.student_model_name, torch_dtype=torch.bfloat16 + ) + teacher_model = AutoModelForImageTextToText.from_pretrained( + cli_args.teacher_model_name, torch_dtype=torch.bfloat16 + ) # Freeze everything except the language model head for name, param in student_model.named_parameters(): From 794a6648e20609605d79da5d5e20c4a57f14786e Mon Sep 17 00:00:00 2001 From: Strongich Date: Sat, 16 May 2026 19:54:05 +0200 Subject: [PATCH 17/39] fix inplace mutation & vlm tests --- tests/experimental/test_gold_trainer.py | 14 +++- trl/experimental/gold/gold_trainer.py | 96 +++++++++++++++---------- 2 files changed, 71 insertions(+), 39 deletions(-) diff --git a/tests/experimental/test_gold_trainer.py b/tests/experimental/test_gold_trainer.py index 3516474433d..752c0c2fc87 100644 --- a/tests/experimental/test_gold_trainer.py +++ b/tests/experimental/test_gold_trainer.py @@ -289,16 +289,20 @@ def qwen3_vl_processor(): @pytest.fixture(scope="module") -def vlm_examples(): +def vlm_dataset(): try: - dataset = load_dataset( + return load_dataset( "trl-internal-testing/zen-image", "conversational_prompt_completion", split="train[:3]", ) except Exception as exc: # pragma: no cover - network/environment dependent pytest.skip(f"zen-image dataset unavailable: {exc}") - return [dict(row) for row in dataset] + + +@pytest.fixture +def vlm_examples(vlm_dataset): + return [dict(row) for row in vlm_dataset] def encode_prompt_completion(tokenizer, prompt, completion): @@ -708,6 +712,7 @@ def test_get_start_and_size_answers_skips_prompt_tokens(): def test_generate_on_policy_outputs_masks_prompt(llama_tokenizer): trainer = GOLDTrainer.__new__(GOLDTrainer) trainer.processing_class = llama_tokenizer + trainer.use_transformers_paged = False prompt_text = "<|begin_of_text|><|start_header_id|>user<|end_header_id|>\nHello?<|eot_id|>" completion_text = "<|start_header_id|>assistant<|end_header_id|>\nHi there!" @@ -759,6 +764,7 @@ def generate(self, input_ids, attention_mask, generation_config, return_dict_in_ def test_generate_on_policy_outputs_masks_prompt_smollm(smollm_tokenizer, openr1_examples): trainer = GOLDTrainer.__new__(GOLDTrainer) trainer.processing_class = smollm_tokenizer + trainer.use_transformers_paged = False collator = DataCollatorForChatML(tokenizer=smollm_tokenizer) batch = collator([openr1_examples[0]]) @@ -1418,6 +1424,7 @@ def __init__(self): self.config = SimpleNamespace( _name_or_path="student", vocab_size=17, vision_config=True, model_type="dummy_vlm" ) + self.config.get_text_config = lambda: self.config self.generation_config = SimpleNamespace(eos_token_id=2) self.name_or_path = "student" @@ -1538,6 +1545,7 @@ def __init__(self): self.config = SimpleNamespace( _name_or_path="student", vocab_size=17, vision_config=True, model_type=student_model_type ) + self.config.get_text_config = lambda: self.config self.generation_config = SimpleNamespace(eos_token_id=2) self.name_or_path = "student" diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index 2b511410242..e6086ec3561 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -18,6 +18,7 @@ from collections import defaultdict, deque from collections.abc import Callable from contextlib import nullcontext +from copy import deepcopy from functools import partial from typing import Any, Optional @@ -1219,6 +1220,23 @@ def _build_sequence_batch( return new_attention_mask, new_labels + @staticmethod + def _get_prompt_sequence_key(inputs: dict[str, torch.Tensor | Any], key: str) -> torch.Tensor: + """Align a sequence-length-dependent key with the left-padded prompt tensor.""" + values = inputs[key] + prompts = inputs["prompts"] + prompt_attention_mask = inputs.get("prompt_attention_mask") + + if prompt_attention_mask is None: + return values[:, : prompts.shape[1]] + + prompt_values = values.new_zeros(prompts.shape) + for i, mask in enumerate(prompt_attention_mask.bool()): + prompt_length = int(mask.sum().item()) + if prompt_length: + prompt_values[i, mask] = values[i, :prompt_length] + return prompt_values + @staticmethod def _get_min_completion_start_from_labels(labels: torch.Tensor) -> int: """Return the earliest completion start, ignoring rows with no completion labels.""" @@ -1379,7 +1397,7 @@ def _generate_non_vllm_for_slices(self, slices: list[dict[str, torch.Tensor | An prompt_seq_len = slice_inputs["prompts"].shape[1] for k in ("token_type_ids", "mm_token_type_ids"): if k in updated_slice: - prompt_part = updated_slice[k][:, :prompt_seq_len] + prompt_part = self._get_prompt_sequence_key(slice_inputs, k) comp_part = torch.zeros( new_input_ids.shape[0], new_seq_len - prompt_seq_len, @@ -1396,6 +1414,7 @@ def _generate_non_vllm_for_slices(self, slices: list[dict[str, torch.Tensor | An def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_indices: list[int]): """On-policy generation from raw VLM examples, preserving PIL images for vLLM.""" device = self.accelerator.device + raw_slices = deepcopy(raw_slices) # Phase 1: Collect prompts, images, and raw examples across all on-policy slices all_prompt_ids = [] @@ -1468,43 +1487,43 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in if not self.use_vllm: # Non-vLLM path: generate per-slice using model.generate - for slice_idx in on_policy_indices: - raw_examples, images, prompts, _ = slice_raw_data[slice_idx] - collated = self._vlm_collator(raw_examples) - collated = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in collated.items()} - with unwrap_model_for_generation( - self.model, self.accelerator, generation_kwargs=self.generation_kwargs - ) as unwrapped_model: + with unwrap_model_for_generation( + self.model, self.accelerator, generation_kwargs=self.generation_kwargs + ) as unwrapped_model: + for slice_idx in on_policy_indices: + raw_examples, images, prompts, _ = slice_raw_data[slice_idx] + collated = self._vlm_collator(raw_examples) + collated = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in collated.items()} result = self.generate_on_policy_outputs( unwrapped_model, collated, self.generation_config, self.pad_token_id ) - new_input_ids, new_attention_mask, new_labels, prompt_texts, completion_texts = result - - updated_slice = dict(collated) - updated_slice["input_ids"] = new_input_ids - updated_slice["attention_mask"] = new_attention_mask - updated_slice["labels"] = new_labels - # Rebuild sequence-length-dependent keys to match new input_ids shape - new_seq_len = new_input_ids.shape[1] - prompt_seq_len = collated["prompts"].shape[1] - for k in ("token_type_ids", "mm_token_type_ids"): - if k in updated_slice: - prompt_part = updated_slice[k][:, :prompt_seq_len] - comp_part = torch.zeros( - new_input_ids.shape[0], - new_seq_len - prompt_seq_len, - dtype=updated_slice[k].dtype, - device=new_input_ids.device, - ) - updated_slice[k] = torch.cat([prompt_part, comp_part], dim=1) - updated_slice["original_prompt_text"] = prompt_texts - updated_slice["original_completion_text"] = completion_texts - if self._teacher_processor is not None: - updated_slice["_raw_images"] = images - updated_slice["_raw_prompts"] = prompts + new_input_ids, new_attention_mask, new_labels, prompt_texts, completion_texts = result + + updated_slice = dict(collated) + updated_slice["input_ids"] = new_input_ids + updated_slice["attention_mask"] = new_attention_mask + updated_slice["labels"] = new_labels + # Rebuild sequence-length-dependent keys to match new input_ids shape + new_seq_len = new_input_ids.shape[1] + prompt_seq_len = collated["prompts"].shape[1] + for k in ("token_type_ids", "mm_token_type_ids"): + if k in updated_slice: + prompt_part = self._get_prompt_sequence_key(collated, k) + comp_part = torch.zeros( + new_input_ids.shape[0], + new_seq_len - prompt_seq_len, + dtype=updated_slice[k].dtype, + device=new_input_ids.device, + ) + updated_slice[k] = torch.cat([prompt_part, comp_part], dim=1) + updated_slice["original_prompt_text"] = prompt_texts + updated_slice["original_completion_text"] = completion_texts + if self._teacher_processor is not None: + updated_slice["_raw_images"] = images + updated_slice["_raw_prompts"] = prompts - self._buffered_inputs[slice_idx] = updated_slice - self._buffered_text_logs[slice_idx] = (prompt_texts, completion_texts) + self._buffered_inputs[slice_idx] = updated_slice + self._buffered_text_logs[slice_idx] = (prompt_texts, completion_texts) return # vLLM path: one batched generate call across all slices @@ -2088,9 +2107,14 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N full.replace(prompt, "", 1) for full, prompt in zip(full_texts, prompt_texts, strict=True) ] + if self._is_vlm and self._teacher_processor is None: + teacher_input_ids = inputs["input_ids"] + teacher_labels = inputs["labels"].clone() + teacher_attention_mask = inputs["attention_mask"] + teacher_prompt_length = self._get_min_completion_start_from_labels(inputs["labels"]) # For cross-architecture VLMs, build teacher inputs with image placeholders by processing # prompts through the teacher's processor with raw images, then appending completions. - if self._teacher_processor is not None and "_raw_images" in inputs: + elif self._teacher_processor is not None and "_raw_images" in inputs: raw_images = inputs["_raw_images"] raw_prompts = inputs["_raw_prompts"] # Apply teacher's chat template to get prompt text with correct image placeholders @@ -2343,7 +2367,7 @@ def generate_on_policy_outputs(self, model, inputs, generation_config, pad_token prompt_seq_len = inputs["prompts"].shape[1] for k in ("token_type_ids", "mm_token_type_ids"): if k in generate_kwargs: - generate_kwargs[k] = generate_kwargs[k][:, :prompt_seq_len] + generate_kwargs[k] = self._get_prompt_sequence_key(inputs, k) generated_outputs = model.generate( input_ids=inputs["prompts"], attention_mask=inputs.get("prompt_attention_mask", None), From 9b072f3b03e2a11dd57b1f0f95b0ad2acd1f4374 Mon Sep 17 00:00:00 2001 From: Strongich Date: Sat, 16 May 2026 20:14:08 +0200 Subject: [PATCH 18/39] remove deprecated paged_attention --- tests/experimental/test_gold_trainer.py | 6 --- trl/experimental/gold/gold_trainer.py | 53 ++++++++----------------- 2 files changed, 16 insertions(+), 43 deletions(-) diff --git a/tests/experimental/test_gold_trainer.py b/tests/experimental/test_gold_trainer.py index 752c0c2fc87..245ec4e5bc9 100644 --- a/tests/experimental/test_gold_trainer.py +++ b/tests/experimental/test_gold_trainer.py @@ -712,7 +712,6 @@ def test_get_start_and_size_answers_skips_prompt_tokens(): def test_generate_on_policy_outputs_masks_prompt(llama_tokenizer): trainer = GOLDTrainer.__new__(GOLDTrainer) trainer.processing_class = llama_tokenizer - trainer.use_transformers_paged = False prompt_text = "<|begin_of_text|><|start_header_id|>user<|end_header_id|>\nHello?<|eot_id|>" completion_text = "<|start_header_id|>assistant<|end_header_id|>\nHi there!" @@ -764,7 +763,6 @@ def generate(self, input_ids, attention_mask, generation_config, return_dict_in_ def test_generate_on_policy_outputs_masks_prompt_smollm(smollm_tokenizer, openr1_examples): trainer = GOLDTrainer.__new__(GOLDTrainer) trainer.processing_class = smollm_tokenizer - trainer.use_transformers_paged = False collator = DataCollatorForChatML(tokenizer=smollm_tokenizer) batch = collator([openr1_examples[0]]) @@ -1108,7 +1106,6 @@ def fake_sft_init( top_p=1.0, seq_kd=False, num_generations=1, - use_transformers_paged=False, max_completion_length=16, top_k=0, log_completions=False, @@ -1392,7 +1389,6 @@ def fake_sft_init( top_p=1.0, seq_kd=False, num_generations=1, - use_transformers_paged=False, max_completion_length=16, top_k=0, log_completions=False, @@ -1493,7 +1489,6 @@ def __init__(self, **kwargs): top_p=1.0, seq_kd=False, num_generations=1, - use_transformers_paged=False, max_completion_length=16, top_k=0, log_completions=False, @@ -1577,7 +1572,6 @@ def _make_vlm_trainer_args(use_vllm=False): top_p=1.0, seq_kd=False, num_generations=1, - use_transformers_paged=False, max_completion_length=16, top_k=0, log_completions=False, diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index e6086ec3561..3a23c740328 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -2340,43 +2340,22 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N def generate_on_policy_outputs(self, model, inputs, generation_config, pad_token_id=None): # Generate output with respect to the prompt only - if self.use_transformers_paged: - previous_attn = self.model.config._attn_implementation - if is_flash_attn_2_available(): - model.config._attn_implementation = "paged_attention" - else: - model.config._attn_implementation = "sdpa_paged" - prompt_mask = inputs.get("prompt_attention_mask") - prompts_tensor = inputs["prompts"] - if prompt_mask is not None: - prompt_sequences = [ - row[mask.bool()].detach().cpu().tolist() - for row, mask in zip(prompts_tensor, prompt_mask, strict=True) - ] - else: - prompt_sequences = [row.detach().cpu().tolist() for row in prompts_tensor] - generated_outputs = model.generate_batch(prompt_sequences, generation_config=generation_config) - model.config._attn_implementation = previous_attn - - completion_ids = [output.generated_tokens for output in generated_outputs.values()] - generated_tokens = torch.stack([torch.tensor(ids, device=model.device) for ids in completion_ids]) - else: - generate_kwargs = {k: inputs[k] for k in self._MULTIMODAL_KEYS if k in inputs} - # Slice sequence-length-dependent keys to prompt-only length (e.g. token_type_ids for Gemma, - # mm_token_type_ids for ERNIE-VL) since model.generate receives prompt-only input_ids - prompt_seq_len = inputs["prompts"].shape[1] - for k in ("token_type_ids", "mm_token_type_ids"): - if k in generate_kwargs: - generate_kwargs[k] = self._get_prompt_sequence_key(inputs, k) - generated_outputs = model.generate( - input_ids=inputs["prompts"], - attention_mask=inputs.get("prompt_attention_mask", None), - generation_config=generation_config, - return_dict_in_generate=True, - **generate_kwargs, - ) - # Get the generated token IDs - generated_tokens = generated_outputs.sequences + generate_kwargs = {k: inputs[k] for k in self._MULTIMODAL_KEYS if k in inputs} + # Slice sequence-length-dependent keys to prompt-only length (e.g. token_type_ids for Gemma, + # mm_token_type_ids for ERNIE-VL) since model.generate receives prompt-only input_ids + prompt_seq_len = inputs["prompts"].shape[1] + for k in ("token_type_ids", "mm_token_type_ids"): + if k in generate_kwargs: + generate_kwargs[k] = self._get_prompt_sequence_key(inputs, k) + generated_outputs = model.generate( + input_ids=inputs["prompts"], + attention_mask=inputs.get("prompt_attention_mask", None), + generation_config=generation_config, + return_dict_in_generate=True, + **generate_kwargs, + ) + # Get the generated token IDs + generated_tokens = generated_outputs.sequences batch_size = generated_tokens.size(0) device = generated_tokens.device From fbd9793cd21ba7f8400b0831728ee85e51a00b25 Mon Sep 17 00:00:00 2001 From: Strongich Date: Sat, 16 May 2026 20:38:58 +0200 Subject: [PATCH 19/39] deepcopy only on-policy slices & remove dead code --- trl/experimental/gold/gold_trainer.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index 3a23c740328..05c5f82ae21 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -1414,7 +1414,6 @@ def _generate_non_vllm_for_slices(self, slices: list[dict[str, torch.Tensor | An def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_indices: list[int]): """On-policy generation from raw VLM examples, preserving PIL images for vLLM.""" device = self.accelerator.device - raw_slices = deepcopy(raw_slices) # Phase 1: Collect prompts, images, and raw examples across all on-policy slices all_prompt_ids = [] @@ -1425,7 +1424,7 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in slice_raw_data = {} # per-slice raw data for non-vLLM path for slice_idx in on_policy_indices: - raw_examples = raw_slices[slice_idx] + raw_examples = deepcopy(raw_slices[slice_idx]) # Extract raw PIL images from examples (like GRPOTrainer) if "images" in raw_examples[0]: @@ -2343,7 +2342,6 @@ def generate_on_policy_outputs(self, model, inputs, generation_config, pad_token generate_kwargs = {k: inputs[k] for k in self._MULTIMODAL_KEYS if k in inputs} # Slice sequence-length-dependent keys to prompt-only length (e.g. token_type_ids for Gemma, # mm_token_type_ids for ERNIE-VL) since model.generate receives prompt-only input_ids - prompt_seq_len = inputs["prompts"].shape[1] for k in ("token_type_ids", "mm_token_type_ids"): if k in generate_kwargs: generate_kwargs[k] = self._get_prompt_sequence_key(inputs, k) From 3b5940e802504c47c37ba05976578f99c183b1d7 Mon Sep 17 00:00:00 2001 From: Strongich Date: Sat, 16 May 2026 21:01:53 +0200 Subject: [PATCH 20/39] emit explicit sequence tensors & fix fragile code in utils --- trl/experimental/gold/gold_trainer.py | 140 +++++++++++++------------- trl/experimental/utils.py | 8 +- 2 files changed, 76 insertions(+), 72 deletions(-) diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index 05c5f82ae21..717388c6a2a 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -2089,82 +2089,82 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N teacher_forward_kwargs = student_forward_kwargs if self.use_uld_loss and self.teacher_tokenizer is not None: - if "original_prompt_text" in inputs and "original_completion_text" in inputs: - prompt_texts = inputs["original_prompt_text"] - completion_texts = inputs["original_completion_text"] - full_texts = [p + c for p, c in zip(prompt_texts, completion_texts, strict=True)] - else: - # Fallback: decode student input_ids (current approach) - # WARNING: This may not work perfectly for cross-tokenizer distillation - full_sequences = inputs["input_ids"] - full_texts = self.processing_class.batch_decode(full_sequences, skip_special_tokens=False) - - # Try to split prompt/completion using original prompt length - prompt_lengths = inputs["prompts"].shape[1] - prompt_texts = self.processing_class.batch_decode(inputs["prompts"], skip_special_tokens=False) - completion_texts = [ - full.replace(prompt, "", 1) for full, prompt in zip(full_texts, prompt_texts, strict=True) - ] - if self._is_vlm and self._teacher_processor is None: teacher_input_ids = inputs["input_ids"] teacher_labels = inputs["labels"].clone() teacher_attention_mask = inputs["attention_mask"] teacher_prompt_length = self._get_min_completion_start_from_labels(inputs["labels"]) - # For cross-architecture VLMs, build teacher inputs with image placeholders by processing - # prompts through the teacher's processor with raw images, then appending completions. - elif self._teacher_processor is not None and "_raw_images" in inputs: - raw_images = inputs["_raw_images"] - raw_prompts = inputs["_raw_prompts"] - # Apply teacher's chat template to get prompt text with correct image placeholders - teacher_prompt_texts = self._teacher_processor.apply_chat_template( - raw_prompts, tokenize=False, add_generation_prompt=True - ) - # Build full text (prompt + completion) and process in one call so all tensors - # (input_ids, attention_mask, mm_token_type_ids, pixel_values, ...) are aligned. - teacher_full_texts = [p + c for p, c in zip(teacher_prompt_texts, completion_texts, strict=True)] - teacher_full_processed = self._teacher_processor( - images=raw_images, - text=teacher_full_texts, - padding=True, - return_tensors="pt", - ) - teacher_input_ids = teacher_full_processed["input_ids"] - teacher_attention_mask = teacher_full_processed["attention_mask"] - # Determine prompt lengths after image token expansion to build labels. - # Derive prompt lengths from total sequence length minus completion length. - # Completions are pure text (no images), so the tokenizer gives exact counts. - # This avoids a second image-processing pass through the teacher processor. - teacher_completion_token_lengths = [ - len(self._teacher_processor.tokenizer(ct, add_special_tokens=False)["input_ids"]) - for ct in completion_texts - ] - total_lengths = teacher_attention_mask.sum(dim=1) - teacher_prompt_token_lengths = [ - int(total_lengths[i].item()) - cl for i, cl in enumerate(teacher_completion_token_lengths) - ] - teacher_labels = teacher_input_ids.clone() - teacher_labels[teacher_attention_mask == 0] = -100 - for i, pl in enumerate(teacher_prompt_token_lengths): - teacher_labels[i, :pl] = -100 - teacher_prompt_length = min(teacher_prompt_token_lengths) - # Override teacher_forward_kwargs with all multimodal keys from teacher processing - teacher_forward_kwargs = { - k: teacher_full_processed[k].to(self.accelerator.device) - for k in self._MULTIMODAL_KEYS - if k in teacher_full_processed - } else: - ( - teacher_input_ids, - teacher_labels, - teacher_attention_mask, - teacher_prompt_length, - ) = build_teacher_inputs_from_texts( - self.teacher_tokenizer, - prompt_texts, - completion_texts, - ) + if "original_prompt_text" in inputs and "original_completion_text" in inputs: + prompt_texts = inputs["original_prompt_text"] + completion_texts = inputs["original_completion_text"] + else: + # Fallback: decode student input_ids (current approach) + # WARNING: This may not work perfectly for cross-tokenizer distillation + full_sequences = inputs["input_ids"] + full_texts = self.processing_class.batch_decode(full_sequences, skip_special_tokens=False) + + # Try to split prompt/completion using original prompt length + prompt_lengths = inputs["prompts"].shape[1] + prompt_texts = self.processing_class.batch_decode(inputs["prompts"], skip_special_tokens=False) + completion_texts = [ + full.replace(prompt, "", 1) for full, prompt in zip(full_texts, prompt_texts, strict=True) + ] + + # For cross-architecture VLMs, build teacher inputs with image placeholders by processing + # prompts through the teacher's processor with raw images, then appending completions. + if self._teacher_processor is not None and "_raw_images" in inputs: + raw_images = inputs["_raw_images"] + raw_prompts = inputs["_raw_prompts"] + # Apply teacher's chat template to get prompt text with correct image placeholders + teacher_prompt_texts = self._teacher_processor.apply_chat_template( + raw_prompts, tokenize=False, add_generation_prompt=True + ) + # Build full text (prompt + completion) and process in one call so all tensors + # (input_ids, attention_mask, mm_token_type_ids, pixel_values, ...) are aligned. + teacher_full_texts = [p + c for p, c in zip(teacher_prompt_texts, completion_texts, strict=True)] + teacher_full_processed = self._teacher_processor( + images=raw_images, + text=teacher_full_texts, + padding=True, + return_tensors="pt", + ) + teacher_input_ids = teacher_full_processed["input_ids"] + teacher_attention_mask = teacher_full_processed["attention_mask"] + # Determine prompt lengths after image token expansion to build labels. + # Derive prompt lengths from total sequence length minus completion length. + # Completions are pure text (no images), so the tokenizer gives exact counts. + # This avoids a second image-processing pass through the teacher processor. + teacher_completion_token_lengths = [ + len(self._teacher_processor.tokenizer(ct, add_special_tokens=False)["input_ids"]) + for ct in completion_texts + ] + total_lengths = teacher_attention_mask.sum(dim=1) + teacher_prompt_token_lengths = [ + int(total_lengths[i].item()) - cl for i, cl in enumerate(teacher_completion_token_lengths) + ] + teacher_labels = teacher_input_ids.clone() + teacher_labels[teacher_attention_mask == 0] = -100 + for i, pl in enumerate(teacher_prompt_token_lengths): + teacher_labels[i, :pl] = -100 + teacher_prompt_length = min(teacher_prompt_token_lengths) + # Override teacher_forward_kwargs with all multimodal keys from teacher processing + teacher_forward_kwargs = { + k: teacher_full_processed[k].to(self.accelerator.device) + for k in self._MULTIMODAL_KEYS + if k in teacher_full_processed + } + else: + ( + teacher_input_ids, + teacher_labels, + teacher_attention_mask, + teacher_prompt_length, + ) = build_teacher_inputs_from_texts( + self.teacher_tokenizer, + prompt_texts, + completion_texts, + ) teacher_input_ids = teacher_input_ids.to(self.accelerator.device) teacher_labels = teacher_labels.to(self.accelerator.device) diff --git a/trl/experimental/utils.py b/trl/experimental/utils.py index 5eff94454db..d7e281645fc 100644 --- a/trl/experimental/utils.py +++ b/trl/experimental/utils.py @@ -435,8 +435,12 @@ def _raw_text_from_messages(messages_or_str: Any) -> str: labels[attention_mask == 0] = -100 labels[completion_mask == 0] = -100 - # Build output with vision keys from processed_prompts (pixel_values, image_grid_thw, etc.) - output = processed_prompts + # Build output with non-sequence vision keys from processed_prompts (pixel_values, image_grid_thw, etc.). + output = { + k: v + for k, v in processed_prompts.items() + if k not in ("input_ids", "attention_mask", "token_type_ids", "mm_token_type_ids") + } output["input_ids"] = input_ids output["attention_mask"] = attention_mask output["labels"] = labels From 01284676c9c5e9a070dcb82c04e64d3a3eb2c9ca Mon Sep 17 00:00:00 2001 From: Strongich Date: Sat, 16 May 2026 22:26:39 +0200 Subject: [PATCH 21/39] add safe .get() to VLMcollator & fix separate tokenization edge case for teacher --- trl/experimental/gold/gold_trainer.py | 90 +++++++++++++++++++-------- trl/experimental/utils.py | 2 +- 2 files changed, 65 insertions(+), 27 deletions(-) diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index 717388c6a2a..5d99488e9dd 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -2120,40 +2120,78 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N teacher_prompt_texts = self._teacher_processor.apply_chat_template( raw_prompts, tokenize=False, add_generation_prompt=True ) - # Build full text (prompt + completion) and process in one call so all tensors - # (input_ids, attention_mask, mm_token_type_ids, pixel_values, ...) are aligned. - teacher_full_texts = [p + c for p, c in zip(teacher_prompt_texts, completion_texts, strict=True)] - teacher_full_processed = self._teacher_processor( + teacher_prompt_processed = self._teacher_processor( images=raw_images, - text=teacher_full_texts, + text=teacher_prompt_texts, padding=True, return_tensors="pt", ) - teacher_input_ids = teacher_full_processed["input_ids"] - teacher_attention_mask = teacher_full_processed["attention_mask"] - # Determine prompt lengths after image token expansion to build labels. - # Derive prompt lengths from total sequence length minus completion length. - # Completions are pure text (no images), so the tokenizer gives exact counts. - # This avoids a second image-processing pass through the teacher processor. - teacher_completion_token_lengths = [ - len(self._teacher_processor.tokenizer(ct, add_special_tokens=False)["input_ids"]) - for ct in completion_texts - ] - total_lengths = teacher_attention_mask.sum(dim=1) - teacher_prompt_token_lengths = [ - int(total_lengths[i].item()) - cl for i, cl in enumerate(teacher_completion_token_lengths) - ] - teacher_labels = teacher_input_ids.clone() - teacher_labels[teacher_attention_mask == 0] = -100 - for i, pl in enumerate(teacher_prompt_token_lengths): - teacher_labels[i, :pl] = -100 + teacher_completion_token_ids = self._teacher_processor.tokenizer( + completion_texts, add_special_tokens=False + )["input_ids"] + + pad_token_id = self.teacher_tokenizer.pad_token_id + eos_token_id = self.teacher_tokenizer.eos_token_id + teacher_sequences = [] + teacher_attention_masks = [] + teacher_labels_list = [] + teacher_prompt_token_lengths = [] + teacher_sequence_kwargs = defaultdict(list) + teacher_sequence_keys = ("token_type_ids", "mm_token_type_ids") + + for row, completion_ids in enumerate(teacher_completion_token_ids): + prompt_mask = teacher_prompt_processed["attention_mask"][row].bool() + prompt_ids = teacher_prompt_processed["input_ids"][row][prompt_mask].tolist() + if eos_token_id is not None and prompt_ids and prompt_ids[-1] == eos_token_id: + prompt_ids = prompt_ids[:-1] + + teacher_prompt_token_lengths.append(len(prompt_ids)) + sequence = list(prompt_ids) + sequence.extend(completion_ids) + if eos_token_id is not None: + sequence.append(eos_token_id) + + seq_tensor = torch.tensor(sequence, dtype=torch.long) + teacher_sequences.append(seq_tensor) + teacher_attention_masks.append(torch.ones_like(seq_tensor)) + + labels = seq_tensor.clone() + labels[: len(prompt_ids)] = -100 + if pad_token_id is not None: + labels[labels == pad_token_id] = -100 + teacher_labels_list.append(labels) + + for key in teacher_sequence_keys: + if key in teacher_prompt_processed: + prompt_values = teacher_prompt_processed[key][row][prompt_mask] + if eos_token_id is not None: + prompt_values = prompt_values[: len(prompt_ids)] + completion_values = torch.zeros( + len(sequence) - len(prompt_ids), + dtype=prompt_values.dtype, + device=prompt_values.device, + ) + teacher_sequence_kwargs[key].append(torch.cat((prompt_values, completion_values))) + + teacher_input_ids = pad( + teacher_sequences, + padding_side="right", + padding_value=pad_token_id if pad_token_id is not None else 0, + ) + teacher_attention_mask = pad(teacher_attention_masks, padding_side="right", padding_value=0).bool() + teacher_labels = pad(teacher_labels_list, padding_side="right", padding_value=-100) teacher_prompt_length = min(teacher_prompt_token_lengths) - # Override teacher_forward_kwargs with all multimodal keys from teacher processing + + # Override teacher_forward_kwargs with multimodal keys from teacher processing. teacher_forward_kwargs = { - k: teacher_full_processed[k].to(self.accelerator.device) + k: teacher_prompt_processed[k].to(self.accelerator.device) for k in self._MULTIMODAL_KEYS - if k in teacher_full_processed + if k in teacher_prompt_processed and k not in teacher_sequence_keys } + for key, values in teacher_sequence_kwargs.items(): + teacher_forward_kwargs[key] = pad(values, padding_side="right", padding_value=0).to( + self.accelerator.device + ) else: ( teacher_input_ids, diff --git a/trl/experimental/utils.py b/trl/experimental/utils.py index d7e281645fc..fab481a7c50 100644 --- a/trl/experimental/utils.py +++ b/trl/experimental/utils.py @@ -385,7 +385,7 @@ def _raw_text_from_messages(messages_or_str: Any) -> str: completion_mask = torch.cat((torch.zeros_like(prompt_mask), completion_mask), dim=1) if "token_type_ids" in processed_prompts: prompt_token_type_ids = processed_prompts["token_type_ids"] - completion_token_type_ids = processed_completions["token_type_ids"] + completion_token_type_ids = processed_completions.get("token_type_ids", torch.zeros_like(completion_ids)) token_type_ids = torch.cat((prompt_token_type_ids, completion_token_type_ids), dim=1) if "mm_token_type_ids" in processed_prompts: prompt_mm_token_type_ids = processed_prompts["mm_token_type_ids"] From 81866b613bf7ae6ce09797593a6a355c3fb49bec Mon Sep 17 00:00:00 2001 From: Strongich Date: Sat, 16 May 2026 23:01:44 +0200 Subject: [PATCH 22/39] remove manual duplication to use RepeatSampler & add args.vllm_tensor_parallel_size to max_num_seqs (for the TP > 1) --- tests/experimental/test_gold_trainer.py | 178 ++++++++++++---- trl/experimental/gold/gold_trainer.py | 256 +++++++++++++++++------- 2 files changed, 323 insertions(+), 111 deletions(-) diff --git a/tests/experimental/test_gold_trainer.py b/tests/experimental/test_gold_trainer.py index 245ec4e5bc9..5cdfe99c618 100644 --- a/tests/experimental/test_gold_trainer.py +++ b/tests/experimental/test_gold_trainer.py @@ -20,9 +20,16 @@ from transformers import AutoProcessor, AutoTokenizer from trl.experimental.gold import gold_trainer as gold_trainer_module -from trl.experimental.gold.gold_trainer import GOLDTrainer, ULDLoss, build_teacher_inputs_from_texts -from trl.experimental.utils import DataCollatorForChatML, DataCollatorForVisionLanguageChatML -from trl.trainer.utils import identity +from trl.experimental.gold.gold_trainer import ( + GOLDTrainer, + ULDLoss, + build_teacher_inputs_from_texts, +) +from trl.experimental.utils import ( + DataCollatorForChatML, + DataCollatorForVisionLanguageChatML, +) +from trl.trainer.utils import RepeatSampler, identity @pytest.fixture(scope="module") @@ -329,7 +336,12 @@ class RecordingTokenizer: pad_token_id = 0 pad_token = "" - def batch_decode(self, sequences, skip_special_tokens=False, clean_up_tokenization_spaces=False): + def batch_decode( + self, + sequences, + skip_special_tokens=False, + clean_up_tokenization_spaces=False, + ): del skip_special_tokens, clean_up_tokenization_spaces return [" ".join(str(token) for token in sequence) for sequence in sequences] @@ -354,8 +366,14 @@ def batch_decode(self, sequences, skip_special_tokens=False, clean_up_tokenizati ) buffered_inputs = trainer._buffered_inputs[0] - assert torch.equal(buffered_inputs["input_ids"], torch.tensor([[0, 11, 31], [21, 22, 41]], dtype=torch.long)) - assert torch.equal(buffered_inputs["attention_mask"], torch.tensor([[0, 1, 1], [1, 1, 1]], dtype=torch.long)) + assert torch.equal( + buffered_inputs["input_ids"], + torch.tensor([[0, 11, 31], [21, 22, 41]], dtype=torch.long), + ) + assert torch.equal( + buffered_inputs["attention_mask"], + torch.tensor([[0, 1, 1], [1, 1, 1]], dtype=torch.long), + ) assert torch.equal(buffered_inputs["labels"], torch.tensor([[-100, -100, 31], [-100, -100, 41]])) @@ -378,7 +396,12 @@ class RecordingTokenizer: pad_token_id = 9 pad_token = "" - def batch_decode(self, sequences, skip_special_tokens=False, clean_up_tokenization_spaces=False): + def batch_decode( + self, + sequences, + skip_special_tokens=False, + clean_up_tokenization_spaces=False, + ): del clean_up_tokenization_spaces decoded = [] token_map = {5: "A", 6: "B", 9: ""} @@ -465,7 +488,12 @@ class RecordingTokenizer: def __init__(self): self.truncation_side = "right" - def batch_decode(self, sequences, skip_special_tokens=False, clean_up_tokenization_spaces=False): + def batch_decode( + self, + sequences, + skip_special_tokens=False, + clean_up_tokenization_spaces=False, + ): del clean_up_tokenization_spaces token_map = {0: "", 5: "A", 6: "B", 13: "", 42: "C"} decoded = [] @@ -510,8 +538,14 @@ def batch_decode(self, sequences, skip_special_tokens=False, clean_up_tokenizati assert trainer.vllm_generation.prompts == [[5, 13, 6]] assert trainer.vllm_generation.sync_calls == 1 assert torch.equal(buffered_inputs["input_ids"], torch.tensor([[5, 13, 6, 42]], dtype=torch.long)) - assert torch.equal(buffered_inputs["attention_mask"], torch.tensor([[1, 1, 1, 1]], dtype=torch.long)) - assert torch.equal(buffered_inputs["labels"], torch.tensor([[-100, -100, -100, 42]], dtype=torch.long)) + assert torch.equal( + buffered_inputs["attention_mask"], + torch.tensor([[1, 1, 1, 1]], dtype=torch.long), + ) + assert torch.equal( + buffered_inputs["labels"], + torch.tensor([[-100, -100, -100, 42]], dtype=torch.long), + ) assert buffered_inputs["original_prompt_text"] == ["A B"] assert buffered_inputs["original_completion_text"] == ["C"] assert trainer._buffered_text_logs[0] == (["A B"], ["C"]) @@ -555,7 +589,14 @@ def fake_sft_init( preprocess_logits_for_metrics=None, peft_config=None, ): - del data_collator, train_dataset, eval_dataset, compute_metrics, callbacks, optimizers + del ( + data_collator, + train_dataset, + eval_dataset, + compute_metrics, + callbacks, + optimizers, + ) del preprocess_logits_for_metrics, peft_config self.model = model self.args = args @@ -753,7 +794,10 @@ def generate(self, input_ids, attention_mask, generation_config, return_dict_in_ padded_prompt_len = prompt_tensor.shape[1] assert torch.all(new_labels[0, :padded_prompt_len] == -100) - assert torch.equal(new_labels[0, padded_prompt_len:], torch.tensor(completion_ids, dtype=torch.long)) + assert torch.equal( + new_labels[0, padded_prompt_len:], + torch.tensor(completion_ids, dtype=torch.long), + ) assert prompt_texts[0] == llama_tokenizer.decode(prompt_ids, skip_special_tokens=False) assert completion_texts[0] == llama_tokenizer.decode(completion_ids, skip_special_tokens=False) @@ -779,7 +823,10 @@ def generate(self, input_ids, attention_mask, generation_config, return_dict_in_ new_ids, new_mask, new_labels, prompt_texts, completion_texts = GOLDTrainer.generate_on_policy_outputs( trainer, DummyModel(), - {"prompts": batch["prompts"], "prompt_attention_mask": batch["prompt_attention_mask"]}, + { + "prompts": batch["prompts"], + "prompt_attention_mask": batch["prompt_attention_mask"], + }, generation_config, pad_id, ) @@ -1069,7 +1116,14 @@ def fake_sft_init( preprocess_logits_for_metrics=None, peft_config=None, ): - del data_collator, train_dataset, eval_dataset, compute_metrics, callbacks, optimizers + del ( + data_collator, + train_dataset, + eval_dataset, + compute_metrics, + callbacks, + optimizers, + ) del preprocess_logits_for_metrics, peft_config self.model = model self.args = args @@ -1355,7 +1409,14 @@ def fake_sft_init( preprocess_logits_for_metrics=None, peft_config=None, ): - del data_collator, train_dataset, eval_dataset, compute_metrics, callbacks, optimizers + del ( + data_collator, + train_dataset, + eval_dataset, + compute_metrics, + callbacks, + optimizers, + ) del preprocess_logits_for_metrics, peft_config self.model = model self.args = args @@ -1412,13 +1473,17 @@ def fake_sft_init( def test_gold_trainer_vlm_vllm_init_uses_identity_collator(monkeypatch): """When a VLM processor is used with lmbda > 0 and use_vllm=True, GOLDTrainer should use the identity collator - and store a _vlm_collator for on-the-fly collation. vLLM should be initialized with max_model_length from args.""" + and store a _vlm_collator for on-the-fly collation. vLLM should be initialized with max_model_length from args. + """ captured = {} class DummyStudentModel: def __init__(self): self.config = SimpleNamespace( - _name_or_path="student", vocab_size=17, vision_config=True, model_type="dummy_vlm" + _name_or_path="student", + vocab_size=17, + vision_config=True, + model_type="dummy_vlm", ) self.config.get_text_config = lambda: self.config self.generation_config = SimpleNamespace(eos_token_id=2) @@ -1538,7 +1603,10 @@ def _make_dummy_vlm_models(student_model_type, teacher_model_type): class DummyStudentModel: def __init__(self): self.config = SimpleNamespace( - _name_or_path="student", vocab_size=17, vision_config=True, model_type=student_model_type + _name_or_path="student", + vocab_size=17, + vision_config=True, + model_type=student_model_type, ) self.config.get_text_config = lambda: self.config self.generation_config = SimpleNamespace(eos_token_id=2) @@ -1546,7 +1614,11 @@ def __init__(self): class DummyTeacherModel: def __init__(self): - self.config = SimpleNamespace(_name_or_path="teacher", vision_config=True, model_type=teacher_model_type) + self.config = SimpleNamespace( + _name_or_path="teacher", + vision_config=True, + model_type=teacher_model_type, + ) self.resized_to = None def resize_token_embeddings(self, vocab_size): @@ -1649,7 +1721,11 @@ def patched_auto_processor(name, **kwargs): return sentinel_processor return real_auto_processor_from_pretrained(name, **kwargs) - monkeypatch.setattr(gold_trainer_module.AutoProcessor, "from_pretrained", staticmethod(patched_auto_processor)) + monkeypatch.setattr( + gold_trainer_module.AutoProcessor, + "from_pretrained", + staticmethod(patched_auto_processor), + ) vision_dataset = Dataset.from_dict({"messages": [["dummy"]], "image": ["fake_image"]}) student, teacher = _make_dummy_vlm_models("smolvlm", "qwen2_5_vl") @@ -1709,7 +1785,11 @@ def patched_auto_processor(name, **kwargs): return sentinel_processor return real_auto_processor_from_pretrained(name, **kwargs) - monkeypatch.setattr(gold_trainer_module.AutoProcessor, "from_pretrained", staticmethod(patched_auto_processor)) + monkeypatch.setattr( + gold_trainer_module.AutoProcessor, + "from_pretrained", + staticmethod(patched_auto_processor), + ) # Monkeypatch AutoTokenizer.from_pretrained for ULD teacher tokenizer loading sentinel_tokenizer = SimpleNamespace(pad_token="", eos_token="") @@ -1720,7 +1800,11 @@ def patched_auto_tokenizer(name, **kwargs): return sentinel_tokenizer return real_auto_tokenizer_from_pretrained(name, **kwargs) - monkeypatch.setattr(gold_trainer_module.AutoTokenizer, "from_pretrained", staticmethod(patched_auto_tokenizer)) + monkeypatch.setattr( + gold_trainer_module.AutoTokenizer, + "from_pretrained", + staticmethod(patched_auto_tokenizer), + ) vision_dataset = Dataset.from_dict({"messages": [["dummy"]], "image": ["fake_image"]}) student, teacher = _make_dummy_vlm_models("smolvlm", "qwen2_5_vl") @@ -1819,14 +1903,13 @@ def fake_sft_init( assert trainer._vlm_collator is not None -def test_on_policy_vlm_vllm_duplicates_prompts_for_num_generations(monkeypatch): - """Regression: in the vLLM path, `_generate_on_policy_vlm_raw` must pre-duplicate prompts - `num_generations` times to satisfy `vllm_generation.generate`'s contract (which returns one completion per input - prompt entry; colocate mode hardcodes `n=1`). Without duplication, the redistribution loop raised `IndexError` for - `num_generations > 1`. +def test_on_policy_vlm_vllm_does_not_duplicate_repeated_sampler_batch(monkeypatch): + """The VLM vLLM path must rely on RepeatSampler for `num_generations` duplication. + + `VLLMGeneration.generate` expects the incoming prompt batch to already contain the repeated prompt entries, + matching the text-only path. Duplicating here again would produce `num_generations ** 2` completions. """ num_generations = 3 - num_prompts_per_slice = 2 num_slices = 2 trainer = GOLDTrainer.__new__(GOLDTrainer) @@ -1885,33 +1968,44 @@ class FakeImage: def __init__(self, tag): self.tag = tag + unique_prompts_per_slice = 2 + unique_examples = [ + {"prompt": [{"role": "user", "content": f"q{i}"}], "image": FakeImage(str(i))} + for i in range(num_slices * unique_prompts_per_slice) + ] + sampler = RepeatSampler( + unique_examples, + mini_repeat_count=num_generations, + batch_size=len(unique_examples), + shuffle=False, + ) + sampled_examples = [unique_examples[i] for i in sampler] raw_slices = [ - [ - {"prompt": [{"role": "user", "content": f"q{slice_idx}_{i}"}], "image": FakeImage(f"{slice_idx}_{i}")} - for i in range(num_prompts_per_slice) - ] - for slice_idx in range(num_slices) + sampled_examples[i : i + unique_prompts_per_slice * num_generations] + for i in range(0, len(sampled_examples), unique_prompts_per_slice * num_generations) ] on_policy_indices = list(range(num_slices)) # Bypass multimodal-message helper; its exact shape is irrelevant to this regression. - monkeypatch.setattr(gold_trainer_module, "prepare_multimodal_messages", lambda prompt, images: prompt) + monkeypatch.setattr( + gold_trainer_module, + "prepare_multimodal_messages", + lambda prompt, images: prompt, + ) trainer._generate_on_policy_vlm_raw(raw_slices, on_policy_indices) - total_unique_prompts = num_slices * num_prompts_per_slice - # The fix: prompts (and images) are pre-duplicated `num_generations` times before - # being handed to vllm_generation.generate. - assert received["n_prompts"] == total_unique_prompts * num_generations - assert received["n_images"] == total_unique_prompts * num_generations + total_sampled_prompts = num_slices * unique_prompts_per_slice * num_generations + assert received["n_prompts"] == total_sampled_prompts + assert received["n_images"] == total_sampled_prompts - # Each slice ends up with `num_prompts_per_slice * num_generations` synthetic examples. + # Each slice ends up with one synthetic example per sampled prompt entry. assert len(collated_per_call) == num_slices for synthetic in collated_per_call: - assert len(synthetic) == num_prompts_per_slice * num_generations + assert len(synthetic) == unique_prompts_per_slice * num_generations # Buffers populated for every on-policy slice without IndexError. for slice_idx in on_policy_indices: assert slice_idx in trainer._buffered_inputs _, completion_texts = trainer._buffered_text_logs[slice_idx] - assert len(completion_texts) == num_prompts_per_slice * num_generations + assert len(completion_texts) == unique_prompts_per_slice * num_generations diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index 5d99488e9dd..29110d6ffb2 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -40,7 +40,12 @@ from transformers.processing_utils import ProcessorMixin from transformers.tokenization_utils_base import PreTrainedTokenizerBase from transformers.trainer_utils import EvalPrediction, seed_worker -from transformers.utils import is_datasets_available, is_liger_kernel_available, is_peft_available, is_rich_available +from transformers.utils import ( + is_datasets_available, + is_liger_kernel_available, + is_peft_available, + is_rich_available, +) from ...data_utils import ( is_conversational, @@ -63,7 +68,12 @@ pad, split_tensor_dict, ) -from ..utils import DataCollatorForChatML, DataCollatorForVisionLanguageChatML, empty_cache, truncate_dataset +from ..utils import ( + DataCollatorForChatML, + DataCollatorForVisionLanguageChatML, + empty_cache, + truncate_dataset, +) from .gold_config import GOLDConfig @@ -223,7 +233,12 @@ def build_teacher_inputs_from_texts( teacher_prompt_length = min(prompt_lengths) if prompt_lengths else 0 - return teacher_input_ids, teacher_labels, teacher_attention_mask, teacher_prompt_length + return ( + teacher_input_ids, + teacher_labels, + teacher_attention_mask, + teacher_prompt_length, + ) class ULDLoss(nn.Module): @@ -231,7 +246,13 @@ class ULDLoss(nn.Module): Universal Logit Distillation Loss. """ - def __init__(self, config: GOLDConfig, student_tokenizer=None, teacher_tokenizer=None, device=None): + def __init__( + self, + config: GOLDConfig, + student_tokenizer=None, + teacher_tokenizer=None, + device=None, + ): super().__init__() self.device = device self.crossentropy_weight = config.uld_crossentropy_weight @@ -261,7 +282,13 @@ def __init__(self, config: GOLDConfig, student_tokenizer=None, teacher_tokenizer self._initialize_vocabulary_mapping() def __call__( - self, student_logits, teacher_logits, student_labels, teacher_labels, student_input_ids, teacher_input_ids + self, + student_logits, + teacher_logits, + student_labels, + teacher_labels, + student_input_ids, + teacher_input_ids, ): """ Compute ULD loss with GKD trainer interface. @@ -289,7 +316,12 @@ def __call__( # Compute distillation loss using ULD approximation distillation_loss = self._compute_distillation_loss( - student_logits, teacher_logits, student_labels, teacher_labels, student_input_ids, teacher_input_ids + student_logits, + teacher_logits, + student_labels, + teacher_labels, + student_input_ids, + teacher_input_ids, ) return crossentropy_loss + distillation_loss @@ -327,7 +359,13 @@ def _initialize_vocabulary_mapping(self): self.mapping_tensor = self.mapping_tensor.to(self.device) def _compute_distillation_loss( - self, student_logits, teacher_logits, student_labels, teacher_labels, student_input_ids, teacher_input_ids + self, + student_logits, + teacher_logits, + student_labels, + teacher_labels, + student_input_ids, + teacher_input_ids, ): """ Compute the Universal Logit Distillation loss with token mapping. @@ -445,7 +483,11 @@ def to_canonical_pieces(tok, ids): prev = "" for k in range(len(ids)): # IMPORTANT: Do NOT skip special tokens - we need to align them too - cur = tok.decode(ids[: k + 1], skip_special_tokens=False, clean_up_tokenization_spaces=False) + cur = tok.decode( + ids[: k + 1], + skip_special_tokens=False, + clean_up_tokenization_spaces=False, + ) # Extract the incremental addition (may include spaces/ZWJ/etc.) pieces.append(cur[len(prev) :]) prev = cur @@ -669,11 +711,13 @@ def _compute_hybrid_uld_loss(self, student_aligned, teacher_aligned): if teacher_unmatched_size < max_unmatched_size: teacher_unmatched_sorted = F.pad( - teacher_unmatched_sorted, (0, max_unmatched_size - teacher_unmatched_size) + teacher_unmatched_sorted, + (0, max_unmatched_size - teacher_unmatched_size), ) if student_unmatched_size < max_unmatched_size: student_unmatched_sorted = F.pad( - student_unmatched_sorted, (0, max_unmatched_size - student_unmatched_size) + student_unmatched_sorted, + (0, max_unmatched_size - student_unmatched_size), ) # L1 loss on sorted unmatched tokens @@ -750,13 +794,15 @@ class GOLDTrainer(SFTTrainer): _paper = { "title": "Unlocking On-Policy Distillation for Any Model Family", # docstyle-ignore - "citation": textwrap.dedent("""\ + "citation": textwrap.dedent( + """\ @misc{patino2025unlocking, title = {{Unlocking On-Policy Distillation for Any Model Family}}, author = {Carlos Miguel Patiño and Kashif Rasul and Quentin Gallouédec and Ben Burtenshaw and Sergio Paniego and Vaibhav Srivastav and Thibaud Frere and Ed Beeching and Lewis Tunstall and Leandro von Werra and Thomas Wolf}, year = 2025, url = {https://huggingface.co/spaces/HuggingFaceH4/general-on-policy-logit-distillation}, - }"""), + }""" + ), } def __init__( @@ -767,15 +813,16 @@ def __init__( data_collator: DataCollator | None = None, # type: ignore train_dataset: Dataset | None = None, eval_dataset: Dataset | dict[str, Dataset] | None = None, - processing_class: PreTrainedTokenizerBase - | BaseImageProcessor - | FeatureExtractionMixin - | ProcessorMixin - | None = None, + processing_class: ( + PreTrainedTokenizerBase | BaseImageProcessor | FeatureExtractionMixin | ProcessorMixin | None + ) = None, compute_metrics: Callable[[EvalPrediction], dict] | None = None, callbacks: list[TrainerCallback] | None = None, - optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), - preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, + optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = ( + None, + None, + ), + preprocess_logits_for_metrics: (Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None) = None, peft_config: Optional["PeftConfig"] = None, ): self.model_name_or_path = model if isinstance(model, str) else model.config._name_or_path @@ -1038,7 +1085,9 @@ def __init__( tensor_parallel_size=args.vllm_tensor_parallel_size, gpu_memory_utilization=args.vllm_gpu_memory_utilization, max_model_length=args.vllm_max_model_length or args.max_length, - max_num_seqs=args.per_device_train_batch_size * args.gradient_accumulation_steps, + max_num_seqs=args.per_device_train_batch_size + * args.gradient_accumulation_steps + * args.vllm_tensor_parallel_size, enable_sleep_mode=args.vllm_enable_sleep_mode, model_impl=args.vllm_model_impl, repetition_penalty=getattr(args, "repetition_penalty", 1.0), @@ -1247,7 +1296,11 @@ def _get_min_completion_start_from_labels(labels: torch.Tensor) -> int: return completion_mask[rows_with_completion].long().argmax(dim=1).min().item() @profiling_decorator - def _fill_buffer(self, generation_batch: dict[str, torch.Tensor | Any] | list[dict], buffer_steps: int): + def _fill_buffer( + self, + generation_batch: dict[str, torch.Tensor | Any] | list[dict], + buffer_steps: int, + ): if self._vlm_collator is not None: # Identity collator path: generation_batch is list[dict] with raw PIL images. # Split into chunks via list slicing, then collate on-the-fly per slice. @@ -1282,15 +1335,17 @@ def _fill_buffer(self, generation_batch: dict[str, torch.Tensor | Any] | list[di ex.get("images") or ([ex["image"]] if "image" in ex else None) for ex in raw_slices[i] ] raw_prompts = [ - prepare_multimodal_messages(ex["prompt"], images=imgs) - if imgs is not None - else ex.get("prompt") + ( + prepare_multimodal_messages(ex["prompt"], images=imgs) + if imgs is not None + else ex.get("prompt") + ) for ex, imgs in zip(raw_slices[i], raw_images, strict=True) ] # Collate raw examples on-the-fly for off-policy slices slice_inputs = self._vlm_collator(raw_slices[i]) slice_inputs = { - k: v.to(self.accelerator.device) if isinstance(v, torch.Tensor) else v + k: (v.to(self.accelerator.device) if isinstance(v, torch.Tensor) else v) for k, v in slice_inputs.items() } # Preserve raw PIL images and prompts for cross-architecture teacher processing @@ -1386,7 +1441,13 @@ def _generate_non_vllm_for_slices(self, slices: list[dict[str, torch.Tensor | An self.generation_config, self.pad_token_id, ) - new_input_ids, new_attention_mask, new_labels, prompt_texts, completion_texts = result + ( + new_input_ids, + new_attention_mask, + new_labels, + prompt_texts, + completion_texts, + ) = result updated_slice = dict(slice_inputs) updated_slice["input_ids"] = new_input_ids @@ -1450,9 +1511,11 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in # copied from GRPOTrainer prompts = [ [ - {**msg, "content": [{"type": "text", "text": msg["content"]}]} - if isinstance(msg.get("content"), str) - else msg + ( + {**msg, "content": [{"type": "text", "text": msg["content"]}]} + if isinstance(msg.get("content"), str) + else msg + ) for msg in prompt ] for prompt in prompts @@ -1494,9 +1557,18 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in collated = self._vlm_collator(raw_examples) collated = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in collated.items()} result = self.generate_on_policy_outputs( - unwrapped_model, collated, self.generation_config, self.pad_token_id + unwrapped_model, + collated, + self.generation_config, + self.pad_token_id, ) - new_input_ids, new_attention_mask, new_labels, prompt_texts, completion_texts = result + ( + new_input_ids, + new_attention_mask, + new_labels, + prompt_texts, + completion_texts, + ) = result updated_slice = dict(collated) updated_slice["input_ids"] = new_input_ids @@ -1522,7 +1594,10 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in updated_slice["_raw_prompts"] = prompts self._buffered_inputs[slice_idx] = updated_slice - self._buffered_text_logs[slice_idx] = (prompt_texts, completion_texts) + self._buffered_text_logs[slice_idx] = ( + prompt_texts, + completion_texts, + ) return # vLLM path: one batched generate call across all slices @@ -1533,17 +1608,12 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in self.vllm_generation.sync_weights() self._last_vllm_sync_step = self.state.global_step - # `vllm_generation.generate` returns one completion per input prompt entry and expects the - # caller to pre-duplicate prompts `num_generations` times (same contract as GRPOTrainer). - # Without this duplication, colocate mode (which hardcodes n=1) yields only `len(prompts)` - # completions while the redistribution below expects `len(prompts) * num_generations`. - dup_prompt_ids = [ids for ids in all_prompt_ids for _ in range(self.num_generations)] if any(img is not None for img in all_images): - generate_images = [img for img in all_images for _ in range(self.num_generations)] + generate_images = all_images else: generate_images = None _, completion_ids, _, _ = self.vllm_generation.generate( - prompts=dup_prompt_ids, + prompts=all_prompt_ids, images=generate_images, num_generations=self.num_generations, ) @@ -1555,12 +1625,15 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in if len(comp_ids) > max_completion_length: comp_ids = comp_ids[:max_completion_length] all_completion_texts.append( - self.processing_class.decode(comp_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False) + self.processing_class.decode( + comp_ids, + skip_special_tokens=True, + clean_up_tokenization_spaces=False, + ) ) - # Redistribute completions to slices. Completions now align 1:1 with the duplicated inputs, - # so `comp_idx` is a single running counter; raw example/prompt/image entries still reference - # the original unique example `i` since those are shared across the `num_generations` copies. + # Redistribute completions to slices. The RepeatSampler has already duplicated examples + # `num_generations` times, so completions align 1:1 with the sampled input entries. slice_completions = {idx: [] for idx in on_policy_indices} slice_raw = {idx: [] for idx in on_policy_indices} slice_images = {idx: [] for idx in on_policy_indices} @@ -1570,14 +1643,13 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in comp_idx = 0 for i, slice_idx in enumerate(local_slice_indices): - for _ in range(self.num_generations): - slice_completions[slice_idx].append(all_completion_texts[comp_idx]) - slice_raw[slice_idx].append(all_raw_examples[i]) - slice_images[slice_idx].append(all_images[i]) - slice_prompts[slice_idx].append(all_prompts[i]) - slice_prompts_text[slice_idx].append(all_prompts_text[i]) - slice_prompts_text_special[slice_idx].append(all_prompts_text_with_special[i]) - comp_idx += 1 + slice_completions[slice_idx].append(all_completion_texts[comp_idx]) + slice_raw[slice_idx].append(all_raw_examples[i]) + slice_images[slice_idx].append(all_images[i]) + slice_prompts[slice_idx].append(all_prompts[i]) + slice_prompts_text[slice_idx].append(all_prompts_text[i]) + slice_prompts_text_special[slice_idx].append(all_prompts_text_with_special[i]) + comp_idx += 1 for slice_idx in on_policy_indices: completion_texts = slice_completions[slice_idx] @@ -1592,7 +1664,10 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in # Wrap as content blocks so VLM chat templates (e.g. SmolVLM) that index # `message.content[0]` can render the synthetic assistant turn. synthetic["completion"] = [ - {"role": "assistant", "content": [{"type": "text", "text": completion_texts[i]}]} + { + "role": "assistant", + "content": [{"type": "text", "text": completion_texts[i]}], + } ] synthetic_examples.append(synthetic) @@ -1607,7 +1682,10 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in collated["_raw_prompts"] = prompts_for_slice self._buffered_inputs[slice_idx] = collated - self._buffered_text_logs[slice_idx] = (slice_prompts_text[slice_idx], completion_texts) + self._buffered_text_logs[slice_idx] = ( + slice_prompts_text[slice_idx], + completion_texts, + ) def _process_completions_to_buffer( self, @@ -1671,7 +1749,11 @@ def _process_completions_to_buffer( padded_completion_ids_list.append(truncated_completion_tensor) completion_ids_for_text.append(truncated_completion_tensor.tolist()) completion_attention_masks.append( - torch.ones(len(truncated_completion_tensor), device=device, dtype=torch.long) + torch.ones( + len(truncated_completion_tensor), + device=device, + dtype=torch.long, + ) ) elif len(completion_tensor) < max_completion_length: padding_needed = max_completion_length - len(completion_tensor) @@ -1691,7 +1773,11 @@ def _process_completions_to_buffer( completion_attention_masks.append( torch.cat( [ - torch.ones(len(completion_tensor), device=device, dtype=torch.long), + torch.ones( + len(completion_tensor), + device=device, + dtype=torch.long, + ), torch.zeros(padding_needed, device=device, dtype=torch.long), ] ) @@ -1735,7 +1821,7 @@ def _process_completions_to_buffer( def _prepare_dataset( self, dataset: Dataset | IterableDataset, - processing_class: PreTrainedTokenizerBase | BaseImageProcessor | FeatureExtractionMixin | ProcessorMixin, + processing_class: (PreTrainedTokenizerBase | BaseImageProcessor | FeatureExtractionMixin | ProcessorMixin), args, packing: bool, formatting_func: Callable[[dict], str] | None, @@ -1760,7 +1846,7 @@ def _prepare_dataset( def _prepare_dataset_with_original_text( self, dataset: Dataset | IterableDataset, - processing_class: PreTrainedTokenizerBase | BaseImageProcessor | FeatureExtractionMixin | ProcessorMixin, + processing_class: (PreTrainedTokenizerBase | BaseImageProcessor | FeatureExtractionMixin | ProcessorMixin), args, packing: bool, formatting_func: Callable[[dict], str] | None, @@ -1791,7 +1877,7 @@ def _func(example): column_names = next(iter(dataset)).keys() dataset = dataset.map( maybe_convert_to_chatml, - remove_columns="conversations" if "conversations" in column_names else None, + remove_columns=("conversations" if "conversations" in column_names else None), **map_kwargs, ) @@ -1811,7 +1897,7 @@ def add_eos(example, eos_token): dataset = dataset.map( add_eos, fn_kwargs={"eos_token": processing_class.eos_token}, - remove_columns="messages" if "messages" in column_names else None, # renamed to "text" + remove_columns=("messages" if "messages" in column_names else None), # renamed to "text" **map_kwargs, ) @@ -1830,7 +1916,9 @@ def tokenize_with_original_text(example, processing_class, dataset_text_field, a if is_conversational(example): prompt_ids = processing_class.apply_chat_template( - example["prompt"], return_dict=False, **example.get("chat_template_kwargs", {}) + example["prompt"], + return_dict=False, + **example.get("chat_template_kwargs", {}), ) prompt_completion_ids = processing_class.apply_chat_template( example["prompt"] + example["completion"], @@ -1907,7 +1995,9 @@ def tokenize_with_original_text(example, processing_class, dataset_text_field, a else: # Fallback: use empty prompt and full text as completion full_text = processing_class.apply_chat_template( - messages, tokenize=False, **example.get("chat_template_kwargs", {}) + messages, + tokenize=False, + **example.get("chat_template_kwargs", {}), ) result["original_prompt_text"] = "" result["original_completion_text"] = full_text @@ -1940,7 +2030,11 @@ def tokenize_with_original_text(example, processing_class, dataset_text_field, a result.update( { "input_ids": tokenized.input_ids, - "attention_mask": getattr(tokenized, "attention_mask", [1] * len(tokenized.input_ids)), + "attention_mask": getattr( + tokenized, + "attention_mask", + [1] * len(tokenized.input_ids), + ), } ) @@ -1963,7 +2057,11 @@ def tokenize_with_original_text(example, processing_class, dataset_text_field, a if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` map_kwargs["desc"] = f"Packing {dataset_name} dataset" - columns_to_keep = ["input_ids", "original_prompt_text", "original_completion_text"] + columns_to_keep = [ + "input_ids", + "original_prompt_text", + "original_completion_text", + ] existing_columns = set(dataset.column_names) columns_to_select = [col for col in columns_to_keep if col in existing_columns] @@ -2042,7 +2140,12 @@ def generalized_jsd_loss( # log(a + b) = log(exp(log(a)) + exp(log(b))) -> for mixture beta = torch.tensor(beta, dtype=student_log_probs.dtype, device=student_log_probs.device) mixture_log_probs = torch.logsumexp( - torch.stack([student_log_probs + torch.log1p(-beta), teacher_log_probs + torch.log(beta)]), + torch.stack( + [ + student_log_probs + torch.log1p(-beta), + teacher_log_probs + torch.log(beta), + ] + ), dim=0, ) @@ -2241,7 +2344,9 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N base_student = unwrapped_student.get_decoder() else: base_student = getattr( - unwrapped_student, getattr(unwrapped_student, "base_model_prefix", "model"), unwrapped_student + unwrapped_student, + getattr(unwrapped_student, "base_model_prefix", "model"), + unwrapped_student, ) student_outputs = base_student( @@ -2257,7 +2362,9 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N base_teacher = unwrapped_teacher.get_decoder() else: base_teacher = getattr( - unwrapped_teacher, getattr(unwrapped_teacher, "base_model_prefix", "model"), unwrapped_teacher + unwrapped_teacher, + getattr(unwrapped_teacher, "base_model_prefix", "model"), + unwrapped_teacher, ) with torch.no_grad(): teacher_outputs = base_teacher( @@ -2277,7 +2384,9 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N labels_mask = inputs["labels"] != -100 masked_input_ids = torch.where( - labels_mask, inputs["input_ids"], torch.full_like(inputs["input_ids"], -100) + labels_mask, + inputs["input_ids"], + torch.full_like(inputs["input_ids"], -100), ) true_labels = masked_input_ids[:, 1:].contiguous().reshape(-1) @@ -2431,7 +2540,13 @@ def generate_on_policy_outputs(self, model, inputs, generation_config, pad_token ) ) - return new_input_ids, new_attention_mask, new_labels, prompt_texts, completion_texts + return ( + new_input_ids, + new_attention_mask, + new_labels, + prompt_texts, + completion_texts, + ) def _get_liger_zero3_lm_head_gather_ctx(self, model: nn.Module): if not self.use_liger_gkd_loss: @@ -2456,7 +2571,10 @@ def _get_liger_zero3_lm_head_gather_ctx(self, model: nn.Module): @profiling_decorator def training_step( - self, model: nn.Module, inputs: dict[str, torch.Tensor | Any], num_items_in_batch: int | None = None + self, + model: nn.Module, + inputs: dict[str, torch.Tensor | Any], + num_items_in_batch: int | None = None, ) -> torch.Tensor: """ Perform a training step for the General Online Logit Distillation (GOLD) model. From dcd557be67de2731b17266dfc31e5a2e0158443d Mon Sep 17 00:00:00 2001 From: Strongich Date: Sun, 17 May 2026 00:02:30 +0200 Subject: [PATCH 23/39] remove deepcopy to save memory & empty from memory processed slices --- tests/experimental/test_gold_trainer.py | 34 +++++++++++++++++++++++++ trl/experimental/gold/gold_trainer.py | 14 +++++++--- 2 files changed, 44 insertions(+), 4 deletions(-) diff --git a/tests/experimental/test_gold_trainer.py b/tests/experimental/test_gold_trainer.py index 5cdfe99c618..bb653d6dcb2 100644 --- a/tests/experimental/test_gold_trainer.py +++ b/tests/experimental/test_gold_trainer.py @@ -2009,3 +2009,37 @@ def __init__(self, tag): assert slice_idx in trainer._buffered_inputs _, completion_texts = trainer._buffered_text_logs[slice_idx] assert len(completion_texts) == unique_prompts_per_slice * num_generations + + +def test_training_step_releases_consumed_buffer_slot(monkeypatch): + """Regression: completed accumulation microbatches should not stay referenced in the rollout buffer.""" + trainer = GOLDTrainer.__new__(GOLDTrainer) + trainer.args = SimpleNamespace(gradient_accumulation_steps=3) + trainer.use_liger_gkd_loss = False + trainer._step = 2 # `_prepare_inputs` already advanced after returning slice 1 + trainer._buffered_on_policy = [False, False, False] + trainer._buffered_text_logs = [None, None, None] + trainer._textual_logs = {"prompt": [], "completion": []} + trainer._on_policy_loss_total = 0.0 + trainer._off_policy_loss_total = 0.0 + trainer._on_policy_step_equiv = 0.0 + trainer._off_policy_step_equiv = 0.0 + sentinel = {"pixel_values": torch.ones(1)} + trainer._buffered_inputs = [{"id": 0}, sentinel, {"id": 2}] + + def fake_training_step(self, model, inputs, num_items_in_batch=None): + assert self._buffered_inputs[1] is sentinel + return torch.tensor(1.0) + + monkeypatch.setattr(gold_trainer_module.SFTTrainer, "training_step", fake_training_step) + + loss = GOLDTrainer.training_step( + trainer, + torch.nn.Linear(1, 1), + sentinel, + num_items_in_batch=None, + ) + + assert loss.item() == 1.0 + assert trainer._buffered_inputs == [{"id": 0}, None, {"id": 2}] + assert trainer._off_policy_loss_total == 1.0 diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index 29110d6ffb2..ca934a66dc6 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -18,7 +18,6 @@ from collections import defaultdict, deque from collections.abc import Callable from contextlib import nullcontext -from copy import deepcopy from functools import partial from typing import Any, Optional @@ -1343,7 +1342,7 @@ def _fill_buffer( for ex, imgs in zip(raw_slices[i], raw_images, strict=True) ] # Collate raw examples on-the-fly for off-policy slices - slice_inputs = self._vlm_collator(raw_slices[i]) + slice_inputs = self._vlm_collator([dict(example) for example in raw_slices[i]]) slice_inputs = { k: (v.to(self.accelerator.device) if isinstance(v, torch.Tensor) else v) for k, v in slice_inputs.items() @@ -1485,7 +1484,7 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in slice_raw_data = {} # per-slice raw data for non-vLLM path for slice_idx in on_policy_indices: - raw_examples = deepcopy(raw_slices[slice_idx]) + raw_examples = raw_slices[slice_idx] # Extract raw PIL images from examples (like GRPOTrainer) if "images" in raw_examples[0]: @@ -1554,7 +1553,7 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in ) as unwrapped_model: for slice_idx in on_policy_indices: raw_examples, images, prompts, _ = slice_raw_data[slice_idx] - collated = self._vlm_collator(raw_examples) + collated = self._vlm_collator([dict(example) for example in raw_examples]) collated = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in collated.items()} result = self.generate_on_policy_outputs( unwrapped_model, @@ -2609,6 +2608,13 @@ def training_step( else: self._off_policy_loss_total += loss_scalar self._off_policy_step_equiv += step_equiv + + if ( + self._buffered_inputs is not None + and isinstance(self._buffered_inputs, list) + and slice_idx < len(self._buffered_inputs) + ): + self._buffered_inputs[slice_idx] = None return loss def log(self, logs: dict[str, float], start_time: float | None = None) -> None: From 19cdc21585eb58fe8f4d4df1d2fa2904748d1dbe Mon Sep 17 00:00:00 2001 From: Strongich Date: Sun, 17 May 2026 00:19:30 +0200 Subject: [PATCH 24/39] reduce peak memory by deferring slice collation --- tests/experimental/test_gold_trainer.py | 128 +++++++++++++++- trl/experimental/gold/gold_trainer.py | 189 ++++++++++++++---------- 2 files changed, 238 insertions(+), 79 deletions(-) diff --git a/tests/experimental/test_gold_trainer.py b/tests/experimental/test_gold_trainer.py index bb653d6dcb2..997b205b339 100644 --- a/tests/experimental/test_gold_trainer.py +++ b/tests/experimental/test_gold_trainer.py @@ -1999,10 +1999,8 @@ def __init__(self, tag): assert received["n_prompts"] == total_sampled_prompts assert received["n_images"] == total_sampled_prompts - # Each slice ends up with one synthetic example per sampled prompt entry. - assert len(collated_per_call) == num_slices - for synthetic in collated_per_call: - assert len(synthetic) == unique_prompts_per_slice * num_generations + # Synthetic VLM examples are stored lazily and are not collated until their slice is consumed. + assert len(collated_per_call) == 0 # Buffers populated for every on-policy slice without IndexError. for slice_idx in on_policy_indices: @@ -2010,6 +2008,128 @@ def __init__(self, tag): _, completion_texts = trainer._buffered_text_logs[slice_idx] assert len(completion_texts) == unique_prompts_per_slice * num_generations + first_slice = trainer._materialize_vlm_slice(trainer._buffered_inputs[0]) + assert first_slice["input_ids"].shape[0] == unique_prompts_per_slice * num_generations + assert len(collated_per_call) == 1 + assert len(collated_per_call[0]) == unique_prompts_per_slice * num_generations + + +def test_off_policy_vlm_collates_only_consumed_slice(monkeypatch): + trainer = GOLDTrainer.__new__(GOLDTrainer) + trainer.accelerator = SimpleNamespace(device=torch.device("cpu"), is_main_process=True) + trainer.args = SimpleNamespace(gradient_accumulation_steps=2) + trainer.lmbda = 0.0 + trainer.use_uld_loss = False + trainer.teacher_tokenizer = None + trainer._teacher_processor = None + trainer._step = 0 + trainer.model = SimpleNamespace(training=True) + collated_per_call = [] + + def stub_collator(examples): + collated_per_call.append(list(examples)) + return {"input_ids": torch.zeros(len(examples), 1, dtype=torch.long)} + + trainer._vlm_collator = stub_collator + monkeypatch.setattr(gold_trainer_module, "broadcast_object_list", lambda values, from_process: values) + + generation_batch = [ + {"prompt": [{"role": "user", "content": "q0"}], "image": object()}, + {"prompt": [{"role": "user", "content": "q1"}], "image": object()}, + {"prompt": [{"role": "user", "content": "q2"}], "image": object()}, + {"prompt": [{"role": "user", "content": "q3"}], "image": object()}, + ] + + first_slice = trainer._prepare_inputs(generation_batch) + + assert len(collated_per_call) == 1 + assert first_slice["input_ids"].shape[0] == 2 + assert "_gold_vlm_lazy_examples" in trainer._buffered_inputs[1] + + second_slice = trainer._prepare_inputs(generation_batch) + + assert len(collated_per_call) == 2 + assert second_slice["input_ids"].shape[0] == 2 + + +def test_on_policy_vlm_without_vllm_collates_only_consumed_slice(monkeypatch): + trainer = GOLDTrainer.__new__(GOLDTrainer) + trainer.accelerator = SimpleNamespace(device=torch.device("cpu"), is_main_process=True) + trainer.args = SimpleNamespace(gradient_accumulation_steps=2) + trainer.use_vllm = False + trainer._teacher_processor = None + trainer._buffered_inputs = [None, None] + trainer._buffered_text_logs = [None, None] + trainer._step = 1 + trainer.generation_kwargs = {} + trainer.generation_config = SimpleNamespace() + trainer.pad_token_id = 0 + collated_per_call = [] + + class StubProcessor: + @staticmethod + def apply_chat_template(conversation, add_generation_prompt, tokenize, return_dict, padding): + return { + "input_ids": [[1, 2] for _ in conversation], + "attention_mask": [[1, 1] for _ in conversation], + } + + @staticmethod + def batch_decode(ids, skip_special_tokens): + return [f"prompt_{i}" for i in range(len(ids))] + + @staticmethod + def decode(ids, skip_special_tokens, clean_up_tokenization_spaces): + return "decoded" + + trainer.processing_class = StubProcessor + + def stub_collator(examples): + collated_per_call.append(list(examples)) + batch_size = len(examples) + return { + "prompts": torch.ones(batch_size, 2, dtype=torch.long), + "prompt_attention_mask": torch.ones(batch_size, 2, dtype=torch.long), + "pixel_values": torch.zeros(batch_size, 3, 2, 2), + } + + trainer._vlm_collator = stub_collator + + class FakeModel: + training = True + + @staticmethod + def generate(input_ids, attention_mask, generation_config, return_dict_in_generate, **kwargs): + completion = torch.full((input_ids.shape[0], 1), 3, dtype=torch.long) + return SimpleNamespace(sequences=torch.cat([input_ids, completion], dim=1)) + + trainer.model = FakeModel() + + monkeypatch.setattr( + gold_trainer_module, + "unwrap_model_for_generation", + lambda *args, **kwargs: gold_trainer_module.nullcontext(args[0]), + ) + monkeypatch.setattr(gold_trainer_module, "prepare_multimodal_messages", lambda prompt, images: prompt) + + raw_slices = [ + [{"prompt": [{"role": "user", "content": "q0"}], "image": object()}], + [{"prompt": [{"role": "user", "content": "q1"}], "image": object()}], + ] + + trainer._generate_on_policy_vlm_raw(raw_slices, [0, 1]) + + assert len(collated_per_call) == 0 + assert "_gold_vlm_on_policy_raw_examples" in trainer._buffered_inputs[0] + assert "_gold_vlm_on_policy_raw_examples" in trainer._buffered_inputs[1] + + consumed_slice = trainer._prepare_inputs(raw_slices) + + assert len(collated_per_call) == 1 + assert consumed_slice["input_ids"].shape == (1, 3) + assert "_gold_vlm_on_policy_raw_examples" in trainer._buffered_inputs[0] + assert "_gold_vlm_on_policy_raw_examples" not in trainer._buffered_inputs[1] + def test_training_step_releases_consumed_buffer_slot(monkeypatch): """Regression: completed accumulation microbatches should not stay referenced in the rollout buffer.""" diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index ca934a66dc6..d46b00d4139 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -1192,9 +1192,93 @@ def _prepare_inputs(self, generation_batch: dict[str, torch.Tensor | Any]) -> di slice_idx = self._step % buffer_steps inputs = self._buffered_inputs[slice_idx] + if isinstance(inputs, dict): + if "_gold_vlm_on_policy_raw_examples" in inputs: + inputs, text_logs = self._generate_on_policy_vlm_slice(inputs) + self._buffered_inputs[slice_idx] = inputs + self._buffered_text_logs[slice_idx] = text_logs + elif "_gold_vlm_lazy_examples" in inputs: + inputs = self._materialize_vlm_slice(inputs) + self._buffered_inputs[slice_idx] = inputs self._step += 1 return inputs + def _generate_on_policy_vlm_slice( + self, pending_slice: dict[str, Any] + ) -> tuple[dict[str, torch.Tensor | Any], tuple[list[str], list[str]]]: + """Generate and collate one non-vLLM on-policy VLM slice immediately before it is consumed.""" + raw_examples = pending_slice["_gold_vlm_on_policy_raw_examples"] + collated = self._vlm_collator([dict(example) for example in raw_examples]) + collated = { + k: (v.to(self.accelerator.device) if isinstance(v, torch.Tensor) else v) for k, v in collated.items() + } + + with unwrap_model_for_generation( + self.model, self.accelerator, generation_kwargs=self.generation_kwargs + ) as unwrapped_model: + ( + new_input_ids, + new_attention_mask, + new_labels, + prompt_texts, + completion_texts, + ) = self.generate_on_policy_outputs( + unwrapped_model, + collated, + self.generation_config, + self.pad_token_id, + ) + + updated_slice = dict(collated) + updated_slice["input_ids"] = new_input_ids + updated_slice["attention_mask"] = new_attention_mask + updated_slice["labels"] = new_labels + # Rebuild sequence-length-dependent keys to match new input_ids shape + new_seq_len = new_input_ids.shape[1] + prompt_seq_len = collated["prompts"].shape[1] + for k in ("token_type_ids", "mm_token_type_ids"): + if k in updated_slice: + prompt_part = self._get_prompt_sequence_key(collated, k) + comp_part = torch.zeros( + new_input_ids.shape[0], + new_seq_len - prompt_seq_len, + dtype=updated_slice[k].dtype, + device=new_input_ids.device, + ) + updated_slice[k] = torch.cat([prompt_part, comp_part], dim=1) + updated_slice["original_prompt_text"] = prompt_texts + updated_slice["original_completion_text"] = completion_texts + if self._teacher_processor is not None: + updated_slice["_raw_images"] = pending_slice["_gold_vlm_raw_images"] + updated_slice["_raw_prompts"] = pending_slice["_gold_vlm_raw_prompts"] + + return updated_slice, (prompt_texts, completion_texts) + + def _materialize_vlm_slice(self, pending_slice: dict[str, Any]) -> dict[str, torch.Tensor | Any]: + """Collate one pending VLM slice immediately before it is consumed.""" + slice_inputs = self._vlm_collator([dict(example) for example in pending_slice["_gold_vlm_lazy_examples"]]) + slice_inputs = { + k: (v.to(self.accelerator.device) if isinstance(v, torch.Tensor) else v) for k, v in slice_inputs.items() + } + + if "_gold_vlm_original_prompt_text" in pending_slice: + slice_inputs["original_prompt_text"] = pending_slice["_gold_vlm_original_prompt_text"] + slice_inputs["original_completion_text"] = pending_slice["_gold_vlm_original_completion_text"] + elif self.use_uld_loss and self.teacher_tokenizer is not None: + slice_inputs = self._ensure_original_text_fields(slice_inputs) + if "original_prompt_text" not in slice_inputs or "original_completion_text" not in slice_inputs: + raise ValueError( + "Off-policy batch missing 'original_prompt_text' or 'original_completion_text' fields. " + "When using ULD loss with cross-tokenizer alignment, datasets must be prepared with " + "_prepare_dataset_with_original_text(). Ensure your dataset includes these fields." + ) + + if self._teacher_processor is not None: + slice_inputs["_raw_images"] = pending_slice["_gold_vlm_raw_images"] + slice_inputs["_raw_prompts"] = pending_slice["_gold_vlm_raw_prompts"] + + return slice_inputs + def _decode_completion_texts_from_labels(self, slice_inputs: dict[str, torch.Tensor | Any]) -> list[str] | None: """Decode completion text from labels when raw text is absent.""" labels = slice_inputs.get("labels") @@ -1341,20 +1425,15 @@ def _fill_buffer( ) for ex, imgs in zip(raw_slices[i], raw_images, strict=True) ] - # Collate raw examples on-the-fly for off-policy slices - slice_inputs = self._vlm_collator([dict(example) for example in raw_slices[i]]) slice_inputs = { - k: (v.to(self.accelerator.device) if isinstance(v, torch.Tensor) else v) - for k, v in slice_inputs.items() + "_gold_vlm_lazy_examples": raw_slices[i], + "_gold_vlm_raw_images": raw_images, + "_gold_vlm_raw_prompts": raw_prompts, } - # Preserve raw PIL images and prompts for cross-architecture teacher processing - if self._teacher_processor is not None: - slice_inputs["_raw_images"] = raw_images - slice_inputs["_raw_prompts"] = raw_prompts else: slice_inputs = slices[i] - if self.use_uld_loss and self.teacher_tokenizer is not None: + if self._vlm_collator is None and self.use_uld_loss and self.teacher_tokenizer is not None: slice_inputs = self._ensure_original_text_fields(slice_inputs) if "original_prompt_text" not in slice_inputs or "original_completion_text" not in slice_inputs: raise ValueError( @@ -1473,8 +1552,6 @@ def _generate_non_vllm_for_slices(self, slices: list[dict[str, torch.Tensor | An def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_indices: list[int]): """On-policy generation from raw VLM examples, preserving PIL images for vLLM.""" - device = self.accelerator.device - # Phase 1: Collect prompts, images, and raw examples across all on-policy slices all_prompt_ids = [] all_images = [] @@ -1520,6 +1597,10 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in for prompt in prompts ] + if not self.use_vllm: + slice_raw_data[slice_idx] = (raw_examples, images, prompts, None) + continue + # Tokenize prompts to get prompt token IDs # TODO: add self.tools support tokenized = self.processing_class.apply_chat_template( @@ -1543,62 +1624,21 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in all_raw_examples.append(example) local_slice_indices.append(slice_idx) - all_prompts_text = self.processing_class.batch_decode(all_prompt_ids, skip_special_tokens=True) - all_prompts_text_with_special = self.processing_class.batch_decode(all_prompt_ids, skip_special_tokens=False) - if not self.use_vllm: - # Non-vLLM path: generate per-slice using model.generate - with unwrap_model_for_generation( - self.model, self.accelerator, generation_kwargs=self.generation_kwargs - ) as unwrapped_model: - for slice_idx in on_policy_indices: - raw_examples, images, prompts, _ = slice_raw_data[slice_idx] - collated = self._vlm_collator([dict(example) for example in raw_examples]) - collated = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in collated.items()} - result = self.generate_on_policy_outputs( - unwrapped_model, - collated, - self.generation_config, - self.pad_token_id, - ) - ( - new_input_ids, - new_attention_mask, - new_labels, - prompt_texts, - completion_texts, - ) = result - - updated_slice = dict(collated) - updated_slice["input_ids"] = new_input_ids - updated_slice["attention_mask"] = new_attention_mask - updated_slice["labels"] = new_labels - # Rebuild sequence-length-dependent keys to match new input_ids shape - new_seq_len = new_input_ids.shape[1] - prompt_seq_len = collated["prompts"].shape[1] - for k in ("token_type_ids", "mm_token_type_ids"): - if k in updated_slice: - prompt_part = self._get_prompt_sequence_key(collated, k) - comp_part = torch.zeros( - new_input_ids.shape[0], - new_seq_len - prompt_seq_len, - dtype=updated_slice[k].dtype, - device=new_input_ids.device, - ) - updated_slice[k] = torch.cat([prompt_part, comp_part], dim=1) - updated_slice["original_prompt_text"] = prompt_texts - updated_slice["original_completion_text"] = completion_texts - if self._teacher_processor is not None: - updated_slice["_raw_images"] = images - updated_slice["_raw_prompts"] = prompts - - self._buffered_inputs[slice_idx] = updated_slice - self._buffered_text_logs[slice_idx] = ( - prompt_texts, - completion_texts, - ) + # Non-vLLM path: local generation needs collated pixel tensors, so defer generation too. + for slice_idx in on_policy_indices: + raw_examples, images, prompts, _ = slice_raw_data[slice_idx] + has_images = images is not None and any(img is not None for img in images) + self._buffered_inputs[slice_idx] = { + "_gold_vlm_on_policy_raw_examples": raw_examples, + "_gold_vlm_raw_images": images if self._teacher_processor is not None and has_images else None, + "_gold_vlm_raw_prompts": prompts if self._teacher_processor is not None else None, + } return + all_prompts_text = self.processing_class.batch_decode(all_prompt_ids, skip_special_tokens=True) + all_prompts_text_with_special = self.processing_class.batch_decode(all_prompt_ids, skip_special_tokens=False) + # vLLM path: one batched generate call across all slices if ( self.state.global_step != self._last_vllm_sync_step @@ -1670,17 +1710,16 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in ] synthetic_examples.append(synthetic) - # Collate synthetic examples to get pixel_values + properly tokenized input_ids/labels - collated = self._vlm_collator(synthetic_examples) - collated = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in collated.items()} - collated["original_prompt_text"] = slice_prompts_text_special[slice_idx] - collated["original_completion_text"] = completion_texts - if self._teacher_processor is not None: - has_images = any(img is not None for img in images_for_slice) - collated["_raw_images"] = images_for_slice if has_images else None - collated["_raw_prompts"] = prompts_for_slice - - self._buffered_inputs[slice_idx] = collated + has_images = any(img is not None for img in images_for_slice) + self._buffered_inputs[slice_idx] = { + "_gold_vlm_lazy_examples": synthetic_examples, + "_gold_vlm_original_prompt_text": slice_prompts_text_special[slice_idx], + "_gold_vlm_original_completion_text": completion_texts, + "_gold_vlm_raw_images": images_for_slice + if self._teacher_processor is not None and has_images + else None, + "_gold_vlm_raw_prompts": prompts_for_slice if self._teacher_processor is not None else None, + } self._buffered_text_logs[slice_idx] = ( slice_prompts_text[slice_idx], completion_texts, From 935414761ba1a2aea90600c51d262647f98fb0a7 Mon Sep 17 00:00:00 2001 From: Strongich Date: Sun, 17 May 2026 10:22:49 +0200 Subject: [PATCH 25/39] store lazy payloads when _teacher_processor is not None --- trl/experimental/gold/gold_trainer.py | 30 ++++++++++++--------------- 1 file changed, 13 insertions(+), 17 deletions(-) diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index d46b00d4139..031ff985ce1 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -1411,8 +1411,7 @@ def _fill_buffer( if self._vlm_collator is not None: # Extract raw images and prompts BEFORE collation, since the collator # mutates examples in place (pops "image", overwrites "prompt"). - raw_images = None - raw_prompts = None + slice_inputs = {"_gold_vlm_lazy_examples": raw_slices[i]} if self._teacher_processor is not None: raw_images = [ ex.get("images") or ([ex["image"]] if "image" in ex else None) for ex in raw_slices[i] @@ -1425,11 +1424,8 @@ def _fill_buffer( ) for ex, imgs in zip(raw_slices[i], raw_images, strict=True) ] - slice_inputs = { - "_gold_vlm_lazy_examples": raw_slices[i], - "_gold_vlm_raw_images": raw_images, - "_gold_vlm_raw_prompts": raw_prompts, - } + slice_inputs["_gold_vlm_raw_images"] = raw_images + slice_inputs["_gold_vlm_raw_prompts"] = raw_prompts else: slice_inputs = slices[i] @@ -1629,11 +1625,11 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in for slice_idx in on_policy_indices: raw_examples, images, prompts, _ = slice_raw_data[slice_idx] has_images = images is not None and any(img is not None for img in images) - self._buffered_inputs[slice_idx] = { - "_gold_vlm_on_policy_raw_examples": raw_examples, - "_gold_vlm_raw_images": images if self._teacher_processor is not None and has_images else None, - "_gold_vlm_raw_prompts": prompts if self._teacher_processor is not None else None, - } + pending_slice = {"_gold_vlm_on_policy_raw_examples": raw_examples} + if self._teacher_processor is not None: + pending_slice["_gold_vlm_raw_images"] = images if has_images else None + pending_slice["_gold_vlm_raw_prompts"] = prompts + self._buffered_inputs[slice_idx] = pending_slice return all_prompts_text = self.processing_class.batch_decode(all_prompt_ids, skip_special_tokens=True) @@ -1711,15 +1707,15 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in synthetic_examples.append(synthetic) has_images = any(img is not None for img in images_for_slice) - self._buffered_inputs[slice_idx] = { + pending_slice = { "_gold_vlm_lazy_examples": synthetic_examples, "_gold_vlm_original_prompt_text": slice_prompts_text_special[slice_idx], "_gold_vlm_original_completion_text": completion_texts, - "_gold_vlm_raw_images": images_for_slice - if self._teacher_processor is not None and has_images - else None, - "_gold_vlm_raw_prompts": prompts_for_slice if self._teacher_processor is not None else None, } + if self._teacher_processor is not None: + pending_slice["_gold_vlm_raw_images"] = images_for_slice if has_images else None + pending_slice["_gold_vlm_raw_prompts"] = prompts_for_slice + self._buffered_inputs[slice_idx] = pending_slice self._buffered_text_logs[slice_idx] = ( slice_prompts_text[slice_idx], completion_texts, From a36800c6176e8294b5c657237428b2a273992891 Mon Sep 17 00:00:00 2001 From: Strongich Date: Sun, 17 May 2026 10:40:01 +0200 Subject: [PATCH 26/39] add an explicit teacher_processor for same family VLMs when use_uld_loss = True --- tests/experimental/test_gold_trainer.py | 85 +++++++++++++++++++++++++ trl/experimental/gold/gold_trainer.py | 25 ++++++-- 2 files changed, 104 insertions(+), 6 deletions(-) diff --git a/tests/experimental/test_gold_trainer.py b/tests/experimental/test_gold_trainer.py index 997b205b339..aad2d3531e3 100644 --- a/tests/experimental/test_gold_trainer.py +++ b/tests/experimental/test_gold_trainer.py @@ -1793,6 +1793,7 @@ def patched_auto_processor(name, **kwargs): # Monkeypatch AutoTokenizer.from_pretrained for ULD teacher tokenizer loading sentinel_tokenizer = SimpleNamespace(pad_token="", eos_token="") + sentinel_processor.tokenizer = sentinel_tokenizer real_auto_tokenizer_from_pretrained = AutoTokenizer.from_pretrained def patched_auto_tokenizer(name, **kwargs): @@ -1903,6 +1904,90 @@ def fake_sft_init( assert trainer._vlm_collator is not None +def test_same_architecture_vlm_with_uld_sets_teacher_processor(monkeypatch): + """ULD VLM distillation should use a teacher processor even when the VLM model_type matches.""" + + def fake_sft_init( + self, + model, + args=None, + data_collator=None, + train_dataset=None, + eval_dataset=None, + processing_class=None, + compute_metrics=None, + callbacks=None, + optimizers=None, + preprocess_logits_for_metrics=None, + peft_config=None, + ): + self.data_collator = data_collator + self.model = model + self.args = args + self.processing_class = processing_class + self.accelerator = SimpleNamespace( + device=torch.device("cpu"), + num_processes=1, + prepare_model=lambda module, evaluation_mode=True: module, + ) + self.is_deepspeed_enabled = False + self.is_fsdp_enabled = False + + monkeypatch.setattr(gold_trainer_module.SFTTrainer, "__init__", fake_sft_init) + + processor = AutoProcessor.from_pretrained("HuggingFaceTB/SmolVLM-256M-Instruct") + if processor.tokenizer.pad_token is None: + processor.tokenizer.pad_token = processor.tokenizer.eos_token + + sentinel_processor = SimpleNamespace(_is_sentinel=True) + real_auto_processor_from_pretrained = AutoProcessor.from_pretrained + + def patched_auto_processor(name, **kwargs): + if name == "teacher": + return sentinel_processor + return real_auto_processor_from_pretrained(name, **kwargs) + + monkeypatch.setattr( + gold_trainer_module.AutoProcessor, + "from_pretrained", + staticmethod(patched_auto_processor), + ) + + sentinel_tokenizer = SimpleNamespace(pad_token="", eos_token="") + sentinel_processor.tokenizer = sentinel_tokenizer + real_auto_tokenizer_from_pretrained = AutoTokenizer.from_pretrained + + def patched_auto_tokenizer(name, **kwargs): + if name == "teacher": + return sentinel_tokenizer + return real_auto_tokenizer_from_pretrained(name, **kwargs) + + monkeypatch.setattr( + gold_trainer_module.AutoTokenizer, + "from_pretrained", + staticmethod(patched_auto_tokenizer), + ) + + vision_dataset = Dataset.from_dict({"messages": [["dummy"]], "image": ["fake_image"]}) + student, teacher = _make_dummy_vlm_models("smolvlm", "smolvlm") + args = _make_vlm_trainer_args() + args.use_uld_loss = True + args.teacher_tokenizer_name_or_path = "teacher" + + trainer = GOLDTrainer( + model=student, + teacher_model=teacher, + args=args, + train_dataset=vision_dataset, + processing_class=processor, + ) + + assert trainer._teacher_processor is sentinel_processor + assert trainer.teacher_tokenizer is sentinel_tokenizer + assert trainer.data_collator is identity + assert trainer._vlm_collator is not None + + def test_on_policy_vlm_vllm_does_not_duplicate_repeated_sampler_batch(monkeypatch): """The VLM vLLM path must rely on RepeatSampler for `num_generations` duplication. diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index 031ff985ce1..549acf6daf0 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -859,12 +859,14 @@ def __init__( # Check for cross-architecture VLM distillation student_model_type = model.config.model_type if not isinstance(model, str) else None teacher_model_type = AutoConfig.from_pretrained(teacher_model).model_type - if student_model_type and teacher_model_type != student_model_type: + is_cross_architecture = student_model_type and teacher_model_type != student_model_type + if is_cross_architecture: warnings.warn( f"Cross-architecture VLM distillation detected: student is '{student_model_type}', " f"teacher is '{teacher_model_type}'. Images will be processed separately through each " "model's processor, which may increase memory usage and computation time." ) + if is_cross_architecture or args.use_uld_loss: self._teacher_processor = teacher_proc elif self._is_vlm and not isinstance(teacher_model, str): # Teacher already instantiated — check if it looks like a VLM by checking for a vision config @@ -877,12 +879,14 @@ def __init__( # Check for cross-architecture VLM distillation student_model_type = model.config.model_type if not isinstance(model, str) else None teacher_model_type = teacher_model.config.model_type - if student_model_type and teacher_model_type != student_model_type: + is_cross_architecture = student_model_type and teacher_model_type != student_model_type + if is_cross_architecture: warnings.warn( f"Cross-architecture VLM distillation detected: student is '{student_model_type}', " f"teacher is '{teacher_model_type}'. Images will be processed separately through each " "model's processor, which may increase memory usage and computation time." ) + if is_cross_architecture or args.use_uld_loss: self._teacher_processor = AutoProcessor.from_pretrained(teacher_model.config._name_or_path) if self._teacher_processor is not None and not args.use_uld_loss: raise ValueError( @@ -963,7 +967,11 @@ def __init__( teacher_model = create_model_from_path(teacher_model, **init_kwargs) self.use_uld_loss = args.use_uld_loss self.teacher_tokenizer = None - if args.use_uld_loss and args.teacher_tokenizer_name_or_path is not None: + if args.use_uld_loss and self._teacher_processor is not None: + self.teacher_tokenizer = self._teacher_processor.tokenizer + if self.teacher_tokenizer.pad_token is None: + self.teacher_tokenizer.pad_token = self.teacher_tokenizer.eos_token + elif args.use_uld_loss and args.teacher_tokenizer_name_or_path is not None: self.teacher_tokenizer = AutoTokenizer.from_pretrained(args.teacher_tokenizer_name_or_path) if not hasattr(self.teacher_tokenizer, "pad_token") or self.teacher_tokenizer.pad_token is None: self.teacher_tokenizer.pad_token = self.teacher_tokenizer.eos_token @@ -2248,9 +2256,14 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N full.replace(prompt, "", 1) for full, prompt in zip(full_texts, prompt_texts, strict=True) ] - # For cross-architecture VLMs, build teacher inputs with image placeholders by processing - # prompts through the teacher's processor with raw images, then appending completions. - if self._teacher_processor is not None and "_raw_images" in inputs: + # For VLMs, build teacher inputs with image placeholders by processing prompts through + # the teacher's processor, then appending completions. + if self._teacher_processor is not None: + if "_raw_images" not in inputs: + raise ValueError( + "VLM ULD loss requires raw images in the batch so teacher inputs can be rendered with " + "the teacher processor. Use GOLD's VLM collator path, which preserves raw images." + ) raw_images = inputs["_raw_images"] raw_prompts = inputs["_raw_prompts"] # Apply teacher's chat template to get prompt text with correct image placeholders From 831b44834170854c82d272fbc4f9f2708f29cad6 Mon Sep 17 00:00:00 2001 From: Strongich Date: Sun, 17 May 2026 11:28:08 +0200 Subject: [PATCH 27/39] store original text not from the decoded text --- tests/experimental/test_gold_trainer.py | 14 +++++++++++++- trl/experimental/gold/gold_trainer.py | 5 ++--- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/tests/experimental/test_gold_trainer.py b/tests/experimental/test_gold_trainer.py index aad2d3531e3..0c9bf2b6764 100644 --- a/tests/experimental/test_gold_trainer.py +++ b/tests/experimental/test_gold_trainer.py @@ -2008,6 +2008,8 @@ def test_on_policy_vlm_vllm_does_not_duplicate_repeated_sampler_batch(monkeypatc trainer._buffered_inputs = {} trainer._buffered_text_logs = {} trainer._teacher_processor = None + trainer.use_uld_loss = False + trainer.teacher_tokenizer = None class StubProcessor: @staticmethod @@ -2045,7 +2047,13 @@ def generate(self, prompts, images, num_generations): def stub_collator(synthetic_examples): collated_per_call.append(list(synthetic_examples)) - return {"input_ids": torch.zeros(len(synthetic_examples), 1, dtype=torch.long)} + return { + "input_ids": torch.zeros(len(synthetic_examples), 1, dtype=torch.long), + "original_prompt_text": [example["prompt"][0]["content"] for example in synthetic_examples], + "original_completion_text": [ + example["completion"][0]["content"][0]["text"] for example in synthetic_examples + ], + } trainer._vlm_collator = stub_collator @@ -2095,6 +2103,8 @@ def __init__(self, tag): first_slice = trainer._materialize_vlm_slice(trainer._buffered_inputs[0]) assert first_slice["input_ids"].shape[0] == unique_prompts_per_slice * num_generations + assert first_slice["original_prompt_text"] == [example["prompt"][0]["content"] for example in collated_per_call[0]] + assert all("<" not in prompt for prompt in first_slice["original_prompt_text"]) assert len(collated_per_call) == 1 assert len(collated_per_call[0]) == unique_prompts_per_slice * num_generations @@ -2176,6 +2186,7 @@ def stub_collator(examples): "prompts": torch.ones(batch_size, 2, dtype=torch.long), "prompt_attention_mask": torch.ones(batch_size, 2, dtype=torch.long), "pixel_values": torch.zeros(batch_size, 3, 2, 2), + "original_prompt_text": [example["prompt"][0]["content"] for example in examples], } trainer._vlm_collator = stub_collator @@ -2212,6 +2223,7 @@ def generate(input_ids, attention_mask, generation_config, return_dict_in_genera assert len(collated_per_call) == 1 assert consumed_slice["input_ids"].shape == (1, 3) + assert consumed_slice["original_prompt_text"] == ["q1"] assert "_gold_vlm_on_policy_raw_examples" in trainer._buffered_inputs[0] assert "_gold_vlm_on_policy_raw_examples" not in trainer._buffered_inputs[1] diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index 549acf6daf0..0b294e867ea 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -1254,7 +1254,8 @@ def _generate_on_policy_vlm_slice( device=new_input_ids.device, ) updated_slice[k] = torch.cat([prompt_part, comp_part], dim=1) - updated_slice["original_prompt_text"] = prompt_texts + if "original_prompt_text" not in updated_slice: + updated_slice["original_prompt_text"] = prompt_texts updated_slice["original_completion_text"] = completion_texts if self._teacher_processor is not None: updated_slice["_raw_images"] = pending_slice["_gold_vlm_raw_images"] @@ -1717,8 +1718,6 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in has_images = any(img is not None for img in images_for_slice) pending_slice = { "_gold_vlm_lazy_examples": synthetic_examples, - "_gold_vlm_original_prompt_text": slice_prompts_text_special[slice_idx], - "_gold_vlm_original_completion_text": completion_texts, } if self._teacher_processor is not None: pending_slice["_gold_vlm_raw_images"] = images_for_slice if has_images else None From b23e663c20c1a2344a9a89f3043258d318473ff0 Mon Sep 17 00:00:00 2001 From: Strongich Date: Sun, 17 May 2026 11:54:18 +0200 Subject: [PATCH 28/39] consistent rendering through teacher processor --- tests/experimental/test_gold_trainer.py | 51 ++++++ trl/experimental/gold/gold_trainer.py | 231 ++++++++++++------------ 2 files changed, 165 insertions(+), 117 deletions(-) diff --git a/tests/experimental/test_gold_trainer.py b/tests/experimental/test_gold_trainer.py index 0c9bf2b6764..4d6114e9289 100644 --- a/tests/experimental/test_gold_trainer.py +++ b/tests/experimental/test_gold_trainer.py @@ -1828,6 +1828,7 @@ def patched_auto_tokenizer(name, **kwargs): # _teacher_processor should be set for cross-architecture assert trainer._teacher_processor is not None assert trainer._teacher_processor is sentinel_processor + assert trainer._is_cross_architecture_vlm is True # A cross-architecture warning should have been emitted cross_arch_warnings = [w for w in caught if "Cross-architecture VLM distillation" in str(w.message)] @@ -1894,6 +1895,7 @@ def fake_sft_init( # _teacher_processor should be None for same architecture (zero overhead) assert trainer._teacher_processor is None + assert trainer._is_cross_architecture_vlm is False # No cross-architecture warning should have been emitted cross_arch_warnings = [w for w in caught if "Cross-architecture VLM distillation" in str(w.message)] @@ -1983,11 +1985,57 @@ def patched_auto_tokenizer(name, **kwargs): ) assert trainer._teacher_processor is sentinel_processor + assert trainer._is_cross_architecture_vlm is False assert trainer.teacher_tokenizer is sentinel_tokenizer assert trainer.data_collator is identity assert trainer._vlm_collator is not None +def test_same_architecture_vlm_uld_preserves_raw_images_for_teacher_processor(monkeypatch): + trainer = GOLDTrainer.__new__(GOLDTrainer) + trainer.accelerator = SimpleNamespace(device=torch.device("cpu"), is_main_process=True) + trainer.args = SimpleNamespace(gradient_accumulation_steps=2) + trainer.lmbda = 0.0 + trainer.use_uld_loss = True + trainer.teacher_tokenizer = SimpleNamespace(pad_token_id=0) + trainer._teacher_processor = object() + trainer._is_cross_architecture_vlm = False + trainer._step = 0 + trainer.model = SimpleNamespace(training=True) + + def stub_collator(examples): + return { + "input_ids": torch.zeros(len(examples), 1, dtype=torch.long), + "original_prompt_text": [example["prompt"][0]["content"] for example in examples], + "original_completion_text": [example["completion"][0]["content"] for example in examples], + } + + trainer._vlm_collator = stub_collator + monkeypatch.setattr(gold_trainer_module, "broadcast_object_list", lambda values, from_process: values) + monkeypatch.setattr(gold_trainer_module, "prepare_multimodal_messages", lambda prompt, images: prompt) + + images = [object(), object()] + generation_batch = [ + { + "prompt": [{"role": "user", "content": "q0"}], + "completion": [{"role": "assistant", "content": "a0"}], + "image": images[0], + }, + { + "prompt": [{"role": "user", "content": "q1"}], + "completion": [{"role": "assistant", "content": "a1"}], + "image": images[1], + }, + ] + + first_slice = trainer._prepare_inputs(generation_batch) + + assert first_slice["_raw_images"] == [[images[0]]] + assert first_slice["_raw_prompts"] == [generation_batch[0]["prompt"]] + assert "_gold_vlm_raw_images" in trainer._buffered_inputs[1] + assert "_gold_vlm_raw_prompts" in trainer._buffered_inputs[1] + + def test_on_policy_vlm_vllm_does_not_duplicate_repeated_sampler_batch(monkeypatch): """The VLM vLLM path must rely on RepeatSampler for `num_generations` duplication. @@ -2008,6 +2056,7 @@ def test_on_policy_vlm_vllm_does_not_duplicate_repeated_sampler_batch(monkeypatc trainer._buffered_inputs = {} trainer._buffered_text_logs = {} trainer._teacher_processor = None + trainer._is_cross_architecture_vlm = False trainer.use_uld_loss = False trainer.teacher_tokenizer = None @@ -2117,6 +2166,7 @@ def test_off_policy_vlm_collates_only_consumed_slice(monkeypatch): trainer.use_uld_loss = False trainer.teacher_tokenizer = None trainer._teacher_processor = None + trainer._is_cross_architecture_vlm = False trainer._step = 0 trainer.model = SimpleNamespace(training=True) collated_per_call = [] @@ -2153,6 +2203,7 @@ def test_on_policy_vlm_without_vllm_collates_only_consumed_slice(monkeypatch): trainer.args = SimpleNamespace(gradient_accumulation_steps=2) trainer.use_vllm = False trainer._teacher_processor = None + trainer._is_cross_architecture_vlm = False trainer._buffered_inputs = [None, None] trainer._buffered_text_logs = [None, None] trainer._step = 1 diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index 0b294e867ea..7498ea91556 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -848,6 +848,7 @@ def __init__( # VLM distillation: only VLM-to-VLM is supported. Both student and teacher must be # VLMs so that both receive images and multimodal inputs. self._teacher_processor = None + self._is_cross_architecture_vlm = False if self._is_vlm and isinstance(teacher_model, str): # Teacher not yet instantiated -- validate it's a VLM teacher_proc = AutoProcessor.from_pretrained(teacher_model) @@ -860,6 +861,7 @@ def __init__( student_model_type = model.config.model_type if not isinstance(model, str) else None teacher_model_type = AutoConfig.from_pretrained(teacher_model).model_type is_cross_architecture = student_model_type and teacher_model_type != student_model_type + self._is_cross_architecture_vlm = is_cross_architecture if is_cross_architecture: warnings.warn( f"Cross-architecture VLM distillation detected: student is '{student_model_type}', " @@ -880,6 +882,7 @@ def __init__( student_model_type = model.config.model_type if not isinstance(model, str) else None teacher_model_type = teacher_model.config.model_type is_cross_architecture = student_model_type and teacher_model_type != student_model_type + self._is_cross_architecture_vlm = is_cross_architecture if is_cross_architecture: warnings.warn( f"Cross-architecture VLM distillation detected: student is '{student_model_type}', " @@ -888,7 +891,7 @@ def __init__( ) if is_cross_architecture or args.use_uld_loss: self._teacher_processor = AutoProcessor.from_pretrained(teacher_model.config._name_or_path) - if self._teacher_processor is not None and not args.use_uld_loss: + if self._is_cross_architecture_vlm and not args.use_uld_loss: raise ValueError( "Cross-architecture VLM distillation (student and teacher have different `model_type`) is not " "supported with the standard JSD loss because the models require different image token formats " @@ -2228,130 +2231,124 @@ def generalized_jsd_loss( def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): # Extract multimodal fields for student forward passes student_forward_kwargs = {k: inputs[k] for k in self._MULTIMODAL_KEYS if k in inputs} - # For same-architecture teacher reuses student vision tensors. - # For cross-architecture VLMs, this gets overridden in the ULD branch below. + # Standard JSD reuses student vision tensors. VLM ULD rebuilds teacher inputs with + # the teacher processor below. teacher_forward_kwargs = student_forward_kwargs if self.use_uld_loss and self.teacher_tokenizer is not None: - if self._is_vlm and self._teacher_processor is None: - teacher_input_ids = inputs["input_ids"] - teacher_labels = inputs["labels"].clone() - teacher_attention_mask = inputs["attention_mask"] - teacher_prompt_length = self._get_min_completion_start_from_labels(inputs["labels"]) + if "original_prompt_text" in inputs and "original_completion_text" in inputs: + prompt_texts = inputs["original_prompt_text"] + completion_texts = inputs["original_completion_text"] else: - if "original_prompt_text" in inputs and "original_completion_text" in inputs: - prompt_texts = inputs["original_prompt_text"] - completion_texts = inputs["original_completion_text"] - else: - # Fallback: decode student input_ids (current approach) - # WARNING: This may not work perfectly for cross-tokenizer distillation - full_sequences = inputs["input_ids"] - full_texts = self.processing_class.batch_decode(full_sequences, skip_special_tokens=False) - - # Try to split prompt/completion using original prompt length - prompt_lengths = inputs["prompts"].shape[1] - prompt_texts = self.processing_class.batch_decode(inputs["prompts"], skip_special_tokens=False) - completion_texts = [ - full.replace(prompt, "", 1) for full, prompt in zip(full_texts, prompt_texts, strict=True) - ] + # Fallback: decode student input_ids (current approach) + # WARNING: This may not work perfectly for cross-tokenizer distillation + full_sequences = inputs["input_ids"] + full_texts = self.processing_class.batch_decode(full_sequences, skip_special_tokens=False) + + # Try to split prompt/completion using original prompt length + prompt_lengths = inputs["prompts"].shape[1] + prompt_texts = self.processing_class.batch_decode(inputs["prompts"], skip_special_tokens=False) + completion_texts = [ + full.replace(prompt, "", 1) for full, prompt in zip(full_texts, prompt_texts, strict=True) + ] - # For VLMs, build teacher inputs with image placeholders by processing prompts through - # the teacher's processor, then appending completions. - if self._teacher_processor is not None: - if "_raw_images" not in inputs: - raise ValueError( - "VLM ULD loss requires raw images in the batch so teacher inputs can be rendered with " - "the teacher processor. Use GOLD's VLM collator path, which preserves raw images." - ) - raw_images = inputs["_raw_images"] - raw_prompts = inputs["_raw_prompts"] - # Apply teacher's chat template to get prompt text with correct image placeholders - teacher_prompt_texts = self._teacher_processor.apply_chat_template( - raw_prompts, tokenize=False, add_generation_prompt=True - ) - teacher_prompt_processed = self._teacher_processor( - images=raw_images, - text=teacher_prompt_texts, - padding=True, - return_tensors="pt", + # For VLMs, build teacher inputs with image placeholders by processing prompts through + # the teacher's processor, then appending completions. + if self._teacher_processor is not None: + if "_raw_images" not in inputs: + raise ValueError( + "VLM ULD loss requires raw images in the batch so teacher inputs can be rendered with " + "the teacher processor. Use GOLD's VLM collator path, which preserves raw images." ) - teacher_completion_token_ids = self._teacher_processor.tokenizer( - completion_texts, add_special_tokens=False - )["input_ids"] - - pad_token_id = self.teacher_tokenizer.pad_token_id - eos_token_id = self.teacher_tokenizer.eos_token_id - teacher_sequences = [] - teacher_attention_masks = [] - teacher_labels_list = [] - teacher_prompt_token_lengths = [] - teacher_sequence_kwargs = defaultdict(list) - teacher_sequence_keys = ("token_type_ids", "mm_token_type_ids") - - for row, completion_ids in enumerate(teacher_completion_token_ids): - prompt_mask = teacher_prompt_processed["attention_mask"][row].bool() - prompt_ids = teacher_prompt_processed["input_ids"][row][prompt_mask].tolist() - if eos_token_id is not None and prompt_ids and prompt_ids[-1] == eos_token_id: - prompt_ids = prompt_ids[:-1] - - teacher_prompt_token_lengths.append(len(prompt_ids)) - sequence = list(prompt_ids) - sequence.extend(completion_ids) - if eos_token_id is not None: - sequence.append(eos_token_id) - - seq_tensor = torch.tensor(sequence, dtype=torch.long) - teacher_sequences.append(seq_tensor) - teacher_attention_masks.append(torch.ones_like(seq_tensor)) - - labels = seq_tensor.clone() - labels[: len(prompt_ids)] = -100 - if pad_token_id is not None: - labels[labels == pad_token_id] = -100 - teacher_labels_list.append(labels) - - for key in teacher_sequence_keys: - if key in teacher_prompt_processed: - prompt_values = teacher_prompt_processed[key][row][prompt_mask] - if eos_token_id is not None: - prompt_values = prompt_values[: len(prompt_ids)] - completion_values = torch.zeros( - len(sequence) - len(prompt_ids), - dtype=prompt_values.dtype, - device=prompt_values.device, - ) - teacher_sequence_kwargs[key].append(torch.cat((prompt_values, completion_values))) + raw_images = inputs["_raw_images"] + raw_prompts = inputs["_raw_prompts"] + # Apply teacher's chat template to get prompt text with correct image placeholders + teacher_prompt_texts = self._teacher_processor.apply_chat_template( + raw_prompts, tokenize=False, add_generation_prompt=True + ) + teacher_prompt_processed = self._teacher_processor( + images=raw_images, + text=teacher_prompt_texts, + padding=True, + return_tensors="pt", + ) + teacher_completion_token_ids = self._teacher_processor.tokenizer( + completion_texts, add_special_tokens=False + )["input_ids"] + + pad_token_id = self.teacher_tokenizer.pad_token_id + eos_token_id = self.teacher_tokenizer.eos_token_id + teacher_sequences = [] + teacher_attention_masks = [] + teacher_labels_list = [] + teacher_prompt_token_lengths = [] + teacher_sequence_kwargs = defaultdict(list) + teacher_sequence_keys = ("token_type_ids", "mm_token_type_ids") + + for row, completion_ids in enumerate(teacher_completion_token_ids): + prompt_mask = teacher_prompt_processed["attention_mask"][row].bool() + prompt_ids = teacher_prompt_processed["input_ids"][row][prompt_mask].tolist() + if eos_token_id is not None and prompt_ids and prompt_ids[-1] == eos_token_id: + prompt_ids = prompt_ids[:-1] + + teacher_prompt_token_lengths.append(len(prompt_ids)) + sequence = list(prompt_ids) + sequence.extend(completion_ids) + if eos_token_id is not None: + sequence.append(eos_token_id) + + seq_tensor = torch.tensor(sequence, dtype=torch.long) + teacher_sequences.append(seq_tensor) + teacher_attention_masks.append(torch.ones_like(seq_tensor)) + + labels = seq_tensor.clone() + labels[: len(prompt_ids)] = -100 + if pad_token_id is not None: + labels[labels == pad_token_id] = -100 + teacher_labels_list.append(labels) + + for key in teacher_sequence_keys: + if key in teacher_prompt_processed: + prompt_values = teacher_prompt_processed[key][row][prompt_mask] + if eos_token_id is not None: + prompt_values = prompt_values[: len(prompt_ids)] + completion_values = torch.zeros( + len(sequence) - len(prompt_ids), + dtype=prompt_values.dtype, + device=prompt_values.device, + ) + teacher_sequence_kwargs[key].append(torch.cat((prompt_values, completion_values))) - teacher_input_ids = pad( - teacher_sequences, - padding_side="right", - padding_value=pad_token_id if pad_token_id is not None else 0, - ) - teacher_attention_mask = pad(teacher_attention_masks, padding_side="right", padding_value=0).bool() - teacher_labels = pad(teacher_labels_list, padding_side="right", padding_value=-100) - teacher_prompt_length = min(teacher_prompt_token_lengths) - - # Override teacher_forward_kwargs with multimodal keys from teacher processing. - teacher_forward_kwargs = { - k: teacher_prompt_processed[k].to(self.accelerator.device) - for k in self._MULTIMODAL_KEYS - if k in teacher_prompt_processed and k not in teacher_sequence_keys - } - for key, values in teacher_sequence_kwargs.items(): - teacher_forward_kwargs[key] = pad(values, padding_side="right", padding_value=0).to( - self.accelerator.device - ) - else: - ( - teacher_input_ids, - teacher_labels, - teacher_attention_mask, - teacher_prompt_length, - ) = build_teacher_inputs_from_texts( - self.teacher_tokenizer, - prompt_texts, - completion_texts, + teacher_input_ids = pad( + teacher_sequences, + padding_side="right", + padding_value=pad_token_id if pad_token_id is not None else 0, + ) + teacher_attention_mask = pad(teacher_attention_masks, padding_side="right", padding_value=0).bool() + teacher_labels = pad(teacher_labels_list, padding_side="right", padding_value=-100) + teacher_prompt_length = min(teacher_prompt_token_lengths) + + # Override teacher_forward_kwargs with multimodal keys from teacher processing. + teacher_forward_kwargs = { + k: teacher_prompt_processed[k].to(self.accelerator.device) + for k in self._MULTIMODAL_KEYS + if k in teacher_prompt_processed and k not in teacher_sequence_keys + } + for key, values in teacher_sequence_kwargs.items(): + teacher_forward_kwargs[key] = pad(values, padding_side="right", padding_value=0).to( + self.accelerator.device ) + else: + ( + teacher_input_ids, + teacher_labels, + teacher_attention_mask, + teacher_prompt_length, + ) = build_teacher_inputs_from_texts( + self.teacher_tokenizer, + prompt_texts, + completion_texts, + ) teacher_input_ids = teacher_input_ids.to(self.accelerator.device) teacher_labels = teacher_labels.to(self.accelerator.device) From 61d1efc42478099a44c3073e94435f267a05de85 Mon Sep 17 00:00:00 2001 From: Strongich Date: Sun, 17 May 2026 12:37:07 +0200 Subject: [PATCH 29/39] use skip_special_tokens = True in non-vLLM on-policy path --- tests/experimental/test_gold_trainer.py | 15 ++++++++++++--- trl/experimental/gold/gold_trainer.py | 14 ++++++++++++-- 2 files changed, 24 insertions(+), 5 deletions(-) diff --git a/tests/experimental/test_gold_trainer.py b/tests/experimental/test_gold_trainer.py index 4d6114e9289..971dbf53846 100644 --- a/tests/experimental/test_gold_trainer.py +++ b/tests/experimental/test_gold_trainer.py @@ -2226,7 +2226,15 @@ def batch_decode(ids, skip_special_tokens): @staticmethod def decode(ids, skip_special_tokens, clean_up_tokenization_spaces): - return "decoded" + tokens = [] + for token_id in ids: + if token_id == 9: + if skip_special_tokens: + continue + tokens.append("") + else: + tokens.append(f"tok{token_id}") + return "".join(tokens) trainer.processing_class = StubProcessor @@ -2247,7 +2255,7 @@ class FakeModel: @staticmethod def generate(input_ids, attention_mask, generation_config, return_dict_in_generate, **kwargs): - completion = torch.full((input_ids.shape[0], 1), 3, dtype=torch.long) + completion = torch.tensor([[3, 9]] * input_ids.shape[0], dtype=torch.long) return SimpleNamespace(sequences=torch.cat([input_ids, completion], dim=1)) trainer.model = FakeModel() @@ -2273,8 +2281,9 @@ def generate(input_ids, attention_mask, generation_config, return_dict_in_genera consumed_slice = trainer._prepare_inputs(raw_slices) assert len(collated_per_call) == 1 - assert consumed_slice["input_ids"].shape == (1, 3) + assert consumed_slice["input_ids"].shape == (1, 4) assert consumed_slice["original_prompt_text"] == ["q1"] + assert consumed_slice["original_completion_text"] == ["tok3"] assert "_gold_vlm_on_policy_raw_examples" in trainer._buffered_inputs[0] assert "_gold_vlm_on_policy_raw_examples" not in trainer._buffered_inputs[1] diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index 7498ea91556..3dd8bfca9db 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -1259,12 +1259,22 @@ def _generate_on_policy_vlm_slice( updated_slice[k] = torch.cat([prompt_part, comp_part], dim=1) if "original_prompt_text" not in updated_slice: updated_slice["original_prompt_text"] = prompt_texts - updated_slice["original_completion_text"] = completion_texts + clean_completion_texts = [] + for input_ids, labels in zip(new_input_ids, new_labels, strict=True): + completion_token_ids = input_ids[labels != -100].tolist() + clean_completion_texts.append( + self.processing_class.decode( + completion_token_ids, + skip_special_tokens=True, + clean_up_tokenization_spaces=False, + ) + ) + updated_slice["original_completion_text"] = clean_completion_texts if self._teacher_processor is not None: updated_slice["_raw_images"] = pending_slice["_gold_vlm_raw_images"] updated_slice["_raw_prompts"] = pending_slice["_gold_vlm_raw_prompts"] - return updated_slice, (prompt_texts, completion_texts) + return updated_slice, (prompt_texts, clean_completion_texts) def _materialize_vlm_slice(self, pending_slice: dict[str, Any]) -> dict[str, torch.Tensor | Any]: """Collate one pending VLM slice immediately before it is consumed.""" From 89704c1f773caab60beab106eed7dc3258bcc6ce Mon Sep 17 00:00:00 2001 From: Strongich Date: Sun, 17 May 2026 12:45:03 +0200 Subject: [PATCH 30/39] remove unused slice_prompts_text_special --- trl/experimental/gold/gold_trainer.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index 3dd8bfca9db..2833c230837 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -1655,8 +1655,6 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in return all_prompts_text = self.processing_class.batch_decode(all_prompt_ids, skip_special_tokens=True) - all_prompts_text_with_special = self.processing_class.batch_decode(all_prompt_ids, skip_special_tokens=False) - # vLLM path: one batched generate call across all slices if ( self.state.global_step != self._last_vllm_sync_step @@ -1696,7 +1694,6 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in slice_images = {idx: [] for idx in on_policy_indices} slice_prompts = {idx: [] for idx in on_policy_indices} slice_prompts_text = {idx: [] for idx in on_policy_indices} - slice_prompts_text_special = {idx: [] for idx in on_policy_indices} comp_idx = 0 for i, slice_idx in enumerate(local_slice_indices): @@ -1705,7 +1702,6 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in slice_images[slice_idx].append(all_images[i]) slice_prompts[slice_idx].append(all_prompts[i]) slice_prompts_text[slice_idx].append(all_prompts_text[i]) - slice_prompts_text_special[slice_idx].append(all_prompts_text_with_special[i]) comp_idx += 1 for slice_idx in on_policy_indices: From f4ee78bfb5ee2649eaa2f259629439b7257357d6 Mon Sep 17 00:00:00 2001 From: Strongich Date: Sun, 17 May 2026 15:23:12 +0200 Subject: [PATCH 31/39] generalize VLM model kwargs --- tests/experimental/test_gold_trainer.py | 32 ++++++++++++++++ trl/experimental/gold/gold_trainer.py | 49 +++++++++++++++++-------- 2 files changed, 65 insertions(+), 16 deletions(-) diff --git a/tests/experimental/test_gold_trainer.py b/tests/experimental/test_gold_trainer.py index 971dbf53846..73b0544eae0 100644 --- a/tests/experimental/test_gold_trainer.py +++ b/tests/experimental/test_gold_trainer.py @@ -2245,6 +2245,7 @@ def stub_collator(examples): "prompts": torch.ones(batch_size, 2, dtype=torch.long), "prompt_attention_mask": torch.ones(batch_size, 2, dtype=torch.long), "pixel_values": torch.zeros(batch_size, 3, 2, 2), + "spatial_shapes": torch.tensor([[2, 2]] * batch_size, dtype=torch.long), "original_prompt_text": [example["prompt"][0]["content"] for example in examples], } @@ -2255,6 +2256,8 @@ class FakeModel: @staticmethod def generate(input_ids, attention_mask, generation_config, return_dict_in_generate, **kwargs): + assert "spatial_shapes" in kwargs + assert torch.equal(kwargs["spatial_shapes"], torch.tensor([[2, 2]], dtype=torch.long)) completion = torch.tensor([[3, 9]] * input_ids.shape[0], dtype=torch.long) return SimpleNamespace(sequences=torch.cat([input_ids, completion], dim=1)) @@ -2282,12 +2285,41 @@ def generate(input_ids, attention_mask, generation_config, return_dict_in_genera assert len(collated_per_call) == 1 assert consumed_slice["input_ids"].shape == (1, 4) + assert torch.equal(consumed_slice["spatial_shapes"], torch.tensor([[2, 2]], dtype=torch.long)) assert consumed_slice["original_prompt_text"] == ["q1"] assert consumed_slice["original_completion_text"] == ["tok3"] assert "_gold_vlm_on_policy_raw_examples" in trainer._buffered_inputs[0] assert "_gold_vlm_on_policy_raw_examples" not in trainer._buffered_inputs[1] +def test_model_forward_kwargs_preserve_processor_tensor_fields(): + trainer = GOLDTrainer.__new__(GOLDTrainer) + inputs = { + "input_ids": torch.ones(1, 2, dtype=torch.long), + "attention_mask": torch.ones(1, 2, dtype=torch.long), + "labels": torch.ones(1, 2, dtype=torch.long), + "prompts": torch.ones(1, 1, dtype=torch.long), + "prompt_attention_mask": torch.ones(1, 1, dtype=torch.long), + "completion_mask": torch.ones(1, 2, dtype=torch.long), + "assistant_masks": torch.ones(1, 2, dtype=torch.long), + "original_prompt_text": ["prompt"], + "_raw_images": [object()], + "pixel_values": torch.zeros(1, 3, 2, 2), + "spatial_shapes": torch.tensor([[2, 2]], dtype=torch.long), + "custom_processor_tensor": torch.tensor([1]), + "token_type_ids": torch.zeros(1, 2, dtype=torch.long), + } + + kwargs = trainer._get_model_forward_kwargs(inputs) + + assert set(kwargs) == {"pixel_values", "spatial_shapes", "custom_processor_tensor", "token_type_ids"} + assert set(trainer._get_model_forward_kwargs(inputs, exclude=("token_type_ids",))) == { + "pixel_values", + "spatial_shapes", + "custom_processor_tensor", + } + + def test_training_step_releases_consumed_buffer_slot(monkeypatch): """Regression: completed accumulation microbatches should not stay referenced in the rollout buffer.""" trainer = GOLDTrainer.__new__(GOLDTrainer) diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index 2833c230837..f137a225db8 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -1130,6 +1130,7 @@ def _set_signature_columns_if_needed(self): "image_position_ids", "pixel_attention_mask", "image_sizes", + "spatial_shapes", "token_type_ids", "mm_token_type_ids", ] @@ -1247,7 +1248,7 @@ def _generate_on_policy_vlm_slice( # Rebuild sequence-length-dependent keys to match new input_ids shape new_seq_len = new_input_ids.shape[1] prompt_seq_len = collated["prompts"].shape[1] - for k in ("token_type_ids", "mm_token_type_ids"): + for k in self._SEQUENCE_KEYS: if k in updated_slice: prompt_part = self._get_prompt_sequence_key(collated, k) comp_part = torch.zeros( @@ -1552,7 +1553,7 @@ def _generate_non_vllm_for_slices(self, slices: list[dict[str, torch.Tensor | An # Rebuild sequence-length-dependent keys to match new input_ids shape new_seq_len = new_input_ids.shape[1] prompt_seq_len = slice_inputs["prompts"].shape[1] - for k in ("token_type_ids", "mm_token_type_ids"): + for k in self._SEQUENCE_KEYS: if k in updated_slice: prompt_part = self._get_prompt_sequence_key(slice_inputs, k) comp_part = torch.zeros( @@ -2224,19 +2225,34 @@ def generalized_jsd_loss( else: return jsd - _MULTIMODAL_KEYS = ( - "pixel_values", - "image_grid_thw", - "image_position_ids", - "pixel_attention_mask", - "image_sizes", - "token_type_ids", - "mm_token_type_ids", + _SEQUENCE_KEYS = ("token_type_ids", "mm_token_type_ids") + _MODEL_INPUT_RESERVED_KEYS = frozenset( + ( + "input_ids", + "attention_mask", + "labels", + "prompts", + "prompt_attention_mask", + "completion_mask", + "assistant_masks", + "original_prompt_text", + "original_completion_text", + ) ) + def _get_model_forward_kwargs( + self, inputs: dict[str, torch.Tensor | Any], exclude: tuple[str, ...] = () + ) -> dict[str, torch.Tensor]: + reserved_keys = self._MODEL_INPUT_RESERVED_KEYS | set(exclude) + return { + k: v + for k, v in inputs.items() + if k not in reserved_keys and not k.startswith("_") and isinstance(v, torch.Tensor) + } + def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): # Extract multimodal fields for student forward passes - student_forward_kwargs = {k: inputs[k] for k in self._MULTIMODAL_KEYS if k in inputs} + student_forward_kwargs = self._get_model_forward_kwargs(inputs) # Standard JSD reuses student vision tensors. VLM ULD rebuilds teacher inputs with # the teacher processor below. teacher_forward_kwargs = student_forward_kwargs @@ -2336,9 +2352,10 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N # Override teacher_forward_kwargs with multimodal keys from teacher processing. teacher_forward_kwargs = { - k: teacher_prompt_processed[k].to(self.accelerator.device) - for k in self._MULTIMODAL_KEYS - if k in teacher_prompt_processed and k not in teacher_sequence_keys + k: v.to(self.accelerator.device) + for k, v in self._get_model_forward_kwargs( + teacher_prompt_processed, exclude=teacher_sequence_keys + ).items() } for key, values in teacher_sequence_kwargs.items(): teacher_forward_kwargs[key] = pad(values, padding_side="right", padding_value=0).to( @@ -2535,10 +2552,10 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N def generate_on_policy_outputs(self, model, inputs, generation_config, pad_token_id=None): # Generate output with respect to the prompt only - generate_kwargs = {k: inputs[k] for k in self._MULTIMODAL_KEYS if k in inputs} + generate_kwargs = self._get_model_forward_kwargs(inputs) # Slice sequence-length-dependent keys to prompt-only length (e.g. token_type_ids for Gemma, # mm_token_type_ids for ERNIE-VL) since model.generate receives prompt-only input_ids - for k in ("token_type_ids", "mm_token_type_ids"): + for k in self._SEQUENCE_KEYS: if k in generate_kwargs: generate_kwargs[k] = self._get_prompt_sequence_key(inputs, k) generated_outputs = model.generate( From 279838896f612465089b4e0b50eaca34ae4347a1 Mon Sep 17 00:00:00 2001 From: Strongich Date: Sun, 17 May 2026 15:56:29 +0200 Subject: [PATCH 32/39] remove dead_code & \n separations for the original_text downstream tasks --- tests/experimental/test_gold_trainer.py | 21 ++++++++++++++++++++- trl/experimental/gold/gold_trainer.py | 5 +---- trl/experimental/utils.py | 14 ++++++++++---- 3 files changed, 31 insertions(+), 9 deletions(-) diff --git a/tests/experimental/test_gold_trainer.py b/tests/experimental/test_gold_trainer.py index 73b0544eae0..030be9a8466 100644 --- a/tests/experimental/test_gold_trainer.py +++ b/tests/experimental/test_gold_trainer.py @@ -1193,6 +1193,19 @@ def _get_assistant_texts(examples): return texts +def _get_prompt_turn_texts(example): + texts = [] + for turn in example["prompt"]: + content = turn["content"] + if isinstance(content, list): + text = "\n".join(part["text"] for part in content if isinstance(part, dict) and "text" in part) + else: + text = content + if text: + texts.append(text) + return texts + + def test_vlm_chatml_collator_preserves_completion_smolvlm(smolvlm_processor, qwen3_vl_processor, vlm_examples): # 2048 to not truncate the completion tokens collator = DataCollatorForVisionLanguageChatML(processor=smolvlm_processor, max_length=2048) @@ -1370,7 +1383,13 @@ def test_vlm_collator_original_text_is_untemplated(smolvlm_processor, vlm_exampl f"original_completion_text leaked student special token {special!r}: {raw_completion!r}" ) - for raw_prompt in batch["original_prompt_text"]: + for raw_prompt, example in zip(batch["original_prompt_text"], vlm_examples, strict=True): + prompt_turn_texts = _get_prompt_turn_texts(example) + for text in prompt_turn_texts: + assert text.strip() in raw_prompt + if len(prompt_turn_texts) > 1: + assert "\n".join(prompt_turn_texts) == raw_prompt + assert "".join(prompt_turn_texts) != raw_prompt for special in student_specials: assert special not in raw_prompt, ( f"original_prompt_text leaked student special token {special!r}: {raw_prompt!r}" diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index f137a225db8..84151af3899 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -1284,10 +1284,7 @@ def _materialize_vlm_slice(self, pending_slice: dict[str, Any]) -> dict[str, tor k: (v.to(self.accelerator.device) if isinstance(v, torch.Tensor) else v) for k, v in slice_inputs.items() } - if "_gold_vlm_original_prompt_text" in pending_slice: - slice_inputs["original_prompt_text"] = pending_slice["_gold_vlm_original_prompt_text"] - slice_inputs["original_completion_text"] = pending_slice["_gold_vlm_original_completion_text"] - elif self.use_uld_loss and self.teacher_tokenizer is not None: + if self.use_uld_loss and self.teacher_tokenizer is not None: slice_inputs = self._ensure_original_text_fields(slice_inputs) if "original_prompt_text" not in slice_inputs or "original_completion_text" not in slice_inputs: raise ValueError( diff --git a/trl/experimental/utils.py b/trl/experimental/utils.py index fab481a7c50..08f2c6b5063 100644 --- a/trl/experimental/utils.py +++ b/trl/experimental/utils.py @@ -332,14 +332,20 @@ def _raw_text_from_messages(messages_or_str: Any) -> str: for turn in messages_or_str: content = turn.get("content", "") if isinstance(content, str): - parts.append(content) + if content: + parts.append(content) continue + turn_parts: list[str] = [] for block in content: if isinstance(block, dict) and block.get("type") == "text": - parts.append(block.get("text", "")) + text = block.get("text", "") + if text: + turn_parts.append(text) elif isinstance(block, str): - parts.append(block) - return "".join(parts) + turn_parts.append(block) + if turn_parts: + parts.append("\n".join(turn_parts)) + return "\n".join(parts) raw_prompt_texts = [_raw_text_from_messages(example["prompt"]) for example in examples] raw_completion_texts = [_raw_text_from_messages(example["completion"]) for example in examples] From d9307bb312d07554b19a7c52a6a60c603103fc00 Mon Sep 17 00:00:00 2001 From: Strongich Date: Mon, 18 May 2026 11:14:36 +0200 Subject: [PATCH 33/39] change examples & comletion tokens, comp_idx dtyle fixes --- examples/scripts/gold_vlm.py | 83 +++++++++++++------------ tests/experimental/test_gold_trainer.py | 5 +- trl/experimental/gold/gold_trainer.py | 42 ++++++++++--- 3 files changed, 79 insertions(+), 51 deletions(-) diff --git a/examples/scripts/gold_vlm.py b/examples/scripts/gold_vlm.py index e2d87e6cd07..c46e95c01ef 100644 --- a/examples/scripts/gold_vlm.py +++ b/examples/scripts/gold_vlm.py @@ -13,26 +13,22 @@ # limitations under the License. """ -GOLD VLM distillation on MMK12. +GOLD VLM distillation on GEOQA_R1V. -# Example 1 — Same-family distillation (SmolVLM-500M → SmolVLM-256M) +# Same-family distillation (Qwen3-VL-8B → Qwen3-VL-2B) # Uses JSD loss. Same architecture and tokenizer, so standard distillation works directly. # vLLM enabled for faster on-policy generation. accelerate launch examples/scripts/gold_vlm.py \ - --student_model_name HuggingFaceTB/SmolVLM-256M-Instruct \ - --teacher_model_name HuggingFaceTB/SmolVLM-500M-Instruct \ - --lmbda 0.5 \ - --use_vllm \ - --vllm_mode colocate - -# Example 2 — Cross-family distillation (Qwen2.5-VL-3B → SmolVLM-256M) -# Different architectures have incompatible tokenizers and image token formats, -# so ULD (Universal Logit Distillation) loss is required to align logits across vocabularies. + --student_model_name Qwen/Qwen3-VL-2B-Instruct \ + --teacher_model_name Qwen/Qwen3-VL-8B-Instruct + +# Cross-family distillation (Qwen3-VL-8B → LFM2.5-VL-1.6B) +# Uses ULD loss for different tokenizers/processors. vLLM is disabled because this path uses local VLM generation. accelerate launch examples/scripts/gold_vlm.py \ - --student_model_name HuggingFaceTB/SmolVLM-256M-Instruct \ - --teacher_model_name Qwen/Qwen2.5-VL-3B-Instruct \ + --student_model_name LiquidAI/LFM2.5-VL-1.6B \ + --teacher_model_name Qwen/Qwen3-VL-8B-Instruct \ --use_uld_loss \ - --lmbda 0.0 + --no-use_vllm """ import argparse @@ -45,15 +41,18 @@ from trl.experimental.gold import GOLDConfig, GOLDTrainer -SYSTEM_PROMPT = ( - "You are a helpful AI Assistant that provides well-reasoned and detailed responses. " - "You first think about the reasoning process as an internal monologue and then provide the user with the answer. " - "Respond in the following format: \n...\n\n\n...\n" -) +SYSTEM_PROMPT = "Answer with a single number followed by the ° symbol." + + +def normalize_solution(solution): + solution = str(solution).replace("", "").replace("", "").strip() + if solution and not solution.endswith("°"): + solution = f"{solution}°" + return solution def make_conversation(example): - """Convert MMK12 row into the chat format expected by TRL VLM trainers.""" + """Convert GEOQA_R1V row into the chat format expected by TRL VLM trainers.""" return { "prompt": [ { @@ -64,14 +63,14 @@ def make_conversation(example): "role": "user", "content": [ {"type": "image"}, - {"type": "text", "text": example["question"]}, + {"type": "text", "text": example["problem"]}, ], }, ], "completion": [ { "role": "assistant", - "content": [{"type": "text", "text": str(example["answer"])}], + "content": [{"type": "text", "text": normalize_solution(example["solution"])}], }, ], "image": example["image"], @@ -93,11 +92,11 @@ def convert_to_rgb(example): if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument("--student_model_name", type=str, default="HuggingFaceTB/SmolVLM-256M-Instruct") - parser.add_argument("--teacher_model_name", type=str, default="HuggingFaceTB/SmolVLM-500M-Instruct") - parser.add_argument("--use_uld_loss", action="store_true") + parser.add_argument("--student_model_name", type=str, default="Qwen/Qwen3-VL-2B-Instruct") + parser.add_argument("--teacher_model_name", type=str, default="Qwen/Qwen3-VL-8B-Instruct") parser.add_argument("--lmbda", type=float, default=0.5) - parser.add_argument("--use_vllm", action="store_true") + parser.add_argument("--use_uld_loss", action=argparse.BooleanOptionalAction, default=False) + parser.add_argument("--use_vllm", action=argparse.BooleanOptionalAction, default=True) parser.add_argument("--vllm_mode", type=str, default="colocate") cli_args = parser.parse_args() @@ -118,18 +117,17 @@ def convert_to_rgb(example): processor = AutoProcessor.from_pretrained(cli_args.student_model_name, padding_side="left") - # toy example to fit small GPUs peft_config = LoraConfig( - r=4, - lora_alpha=8, + r=16, + lora_alpha=32, lora_dropout=0.05, - target_modules=["q_proj"], + target_modules=r"^.*language_model.*\.(q_proj|k_proj|v_proj)$", ) # ────────────────────────────────────────────── # Dataset # ────────────────────────────────────────────── - dataset = load_dataset("FanqingM/MMK12", split="train[:5%]") + dataset = load_dataset("leonardPKU/GEOQA_R1V_Train_8K", split="train") dataset = dataset.filter(filter_big_images) dataset = dataset.map(convert_to_rgb) dataset = dataset.map(make_conversation) @@ -138,35 +136,38 @@ def convert_to_rgb(example): # Training config # ────────────────────────────────────────────── args = GOLDConfig( - output_dir="gold-vlm-distillation", + output_dir=( + "gold-vlm-distillation-different-family" if cli_args.use_uld_loss else "gold-vlm-distillation-same-family" + ), # GOLD-specific lmbda=cli_args.lmbda, beta=0.5, - temperature=0.9, - max_completion_length=256, + temperature=0.6, + max_completion_length=128, + max_grad_norm=1.0, teacher_model_name_or_path=cli_args.teacher_model_name, num_generations=1, use_uld_loss=cli_args.use_uld_loss, + uld_crossentropy_weight=0.5, + uld_distillation_weight=0.5, # vLLM use_vllm=cli_args.use_vllm, vllm_mode=cli_args.vllm_mode, vllm_gpu_memory_utilization=0.5, - vllm_max_model_length=8192, - # VLM image tokens expand during processing, so the default max_length (1024) is often too small. - # Which will lead to shifted_student_logits become an empty Tensor. + vllm_max_model_length=1024, max_length=2048, # Training schedule per_device_train_batch_size=2, gradient_accumulation_steps=4, - max_steps=100, - learning_rate=2e-5, + max_steps=300, + learning_rate=1e-4, warmup_steps=10, # Precision bf16=True, # Logging - logging_steps=1, + logging_steps=10, log_completions=True, - report_to="none", + report_to="wandb", ) # ────────────────────────────────────────────── diff --git a/tests/experimental/test_gold_trainer.py b/tests/experimental/test_gold_trainer.py index 030be9a8466..8facd067f60 100644 --- a/tests/experimental/test_gold_trainer.py +++ b/tests/experimental/test_gold_trainer.py @@ -2259,6 +2259,7 @@ def decode(ids, skip_special_tokens, clean_up_tokenization_spaces): def stub_collator(examples): collated_per_call.append(list(examples)) + assert all(example.get("completion") == "" for example in examples) batch_size = len(examples) return { "prompts": torch.ones(batch_size, 2, dtype=torch.long), @@ -2290,8 +2291,8 @@ def generate(input_ids, attention_mask, generation_config, return_dict_in_genera monkeypatch.setattr(gold_trainer_module, "prepare_multimodal_messages", lambda prompt, images: prompt) raw_slices = [ - [{"prompt": [{"role": "user", "content": "q0"}], "image": object()}], - [{"prompt": [{"role": "user", "content": "q1"}], "image": object()}], + [{"prompt": [{"role": "user", "content": "q0"}], "completion": "gold0", "image": object()}], + [{"prompt": [{"role": "user", "content": "q1"}], "completion": "gold1", "image": object()}], ] trainer._generate_on_policy_vlm_raw(raw_slices, [0, 1]) diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index 84151af3899..e8ec71dbc7a 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -1220,7 +1220,17 @@ def _generate_on_policy_vlm_slice( ) -> tuple[dict[str, torch.Tensor | Any], tuple[list[str], list[str]]]: """Generate and collate one non-vLLM on-policy VLM slice immediately before it is consumed.""" raw_examples = pending_slice["_gold_vlm_on_policy_raw_examples"] - collated = self._vlm_collator([dict(example) for example in raw_examples]) + generation_examples = [] + for example in raw_examples: + generation_example = dict(example) + completion = generation_example.get("completion") + generation_example["completion"] = ( + [{"role": "assistant", "content": [{"type": "text", "text": ""}]}] + if isinstance(completion, list) + else "" + ) + generation_examples.append(generation_example) + collated = self._vlm_collator(generation_examples) collated = { k: (v.to(self.accelerator.device) if isinstance(v, torch.Tensor) else v) for k, v in collated.items() } @@ -1250,11 +1260,12 @@ def _generate_on_policy_vlm_slice( prompt_seq_len = collated["prompts"].shape[1] for k in self._SEQUENCE_KEYS: if k in updated_slice: + sequence_dtype = updated_slice[k].dtype prompt_part = self._get_prompt_sequence_key(collated, k) comp_part = torch.zeros( new_input_ids.shape[0], new_seq_len - prompt_seq_len, - dtype=updated_slice[k].dtype, + dtype=sequence_dtype, device=new_input_ids.device, ) updated_slice[k] = torch.cat([prompt_part, comp_part], dim=1) @@ -1624,7 +1635,7 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in add_generation_prompt=True, tokenize=True, return_dict=True, - padding=True, + processor_kwargs={"padding": True}, ) prompt_ids_list = [ [tok for tok, m in zip(ids, mask, strict=True) if m] @@ -1693,14 +1704,12 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in slice_prompts = {idx: [] for idx in on_policy_indices} slice_prompts_text = {idx: [] for idx in on_policy_indices} - comp_idx = 0 for i, slice_idx in enumerate(local_slice_indices): - slice_completions[slice_idx].append(all_completion_texts[comp_idx]) + slice_completions[slice_idx].append(all_completion_texts[i]) slice_raw[slice_idx].append(all_raw_examples[i]) slice_images[slice_idx].append(all_images[i]) slice_prompts[slice_idx].append(all_prompts[i]) slice_prompts_text[slice_idx].append(all_prompts_text[i]) - comp_idx += 1 for slice_idx in on_policy_indices: completion_texts = slice_completions[slice_idx] @@ -2285,6 +2294,23 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N teacher_prompt_texts = self._teacher_processor.apply_chat_template( raw_prompts, tokenize=False, add_generation_prompt=True ) + teacher_completion_texts = [] + for raw_prompt, completion_text, teacher_prompt_text in zip( + raw_prompts, completion_texts, teacher_prompt_texts, strict=True + ): + teacher_completion = [ + { + "role": "assistant", + "content": [{"type": "text", "text": completion_text}], + } + ] + teacher_prompt_completion_text = self._teacher_processor.apply_chat_template( + raw_prompt + teacher_completion, tokenize=False, add_generation_prompt=False + ) + if teacher_prompt_completion_text.startswith(teacher_prompt_text): + teacher_completion_texts.append(teacher_prompt_completion_text[len(teacher_prompt_text) :]) + else: + teacher_completion_texts.append(completion_text) teacher_prompt_processed = self._teacher_processor( images=raw_images, text=teacher_prompt_texts, @@ -2292,7 +2318,7 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N return_tensors="pt", ) teacher_completion_token_ids = self._teacher_processor.tokenizer( - completion_texts, add_special_tokens=False + teacher_completion_texts, add_special_tokens=False )["input_ids"] pad_token_id = self.teacher_tokenizer.pad_token_id @@ -2313,7 +2339,7 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N teacher_prompt_token_lengths.append(len(prompt_ids)) sequence = list(prompt_ids) sequence.extend(completion_ids) - if eos_token_id is not None: + if eos_token_id is not None and eos_token_id not in completion_ids: sequence.append(eos_token_id) seq_tensor = torch.tensor(sequence, dtype=torch.long) From 81b606384b370b6b42606e2b77b471c13a7eae51 Mon Sep 17 00:00:00 2001 From: Strongich Date: Sat, 23 May 2026 11:57:44 +0200 Subject: [PATCH 34/39] delete dead code & fix eval collator --- tests/experimental/test_gold_trainer.py | 65 +++++++++++++++++++++++++ trl/experimental/gold/gold_trainer.py | 40 +++++++++------ 2 files changed, 89 insertions(+), 16 deletions(-) diff --git a/tests/experimental/test_gold_trainer.py b/tests/experimental/test_gold_trainer.py index 8facd067f60..1cf6a54403f 100644 --- a/tests/experimental/test_gold_trainer.py +++ b/tests/experimental/test_gold_trainer.py @@ -2216,6 +2216,71 @@ def stub_collator(examples): assert second_slice["input_ids"].shape[0] == 2 +def test_eval_vlm_collates_raw_batch_off_policy(): + """During eval the identity collator yields raw dicts; `_prepare_inputs` must collate them off-policy. + + Regression test: previously the eval branch returned the raw `list[dict]` unchanged, which crashed `compute_loss` + on `inputs["input_ids"]`. Eval must run the VLM collator over the whole batch (no slicing, no buffering, no + generation). + """ + trainer = GOLDTrainer.__new__(GOLDTrainer) + trainer.accelerator = SimpleNamespace(device=torch.device("cpu"), is_main_process=True) + trainer.model = SimpleNamespace(training=False) + trainer.use_uld_loss = False + trainer.teacher_tokenizer = None + trainer._teacher_processor = None + collated_per_call = [] + + def stub_collator(examples): + collated_per_call.append(list(examples)) + return {"input_ids": torch.zeros(len(examples), 1, dtype=torch.long)} + + trainer._vlm_collator = stub_collator + + generation_batch = [ + {"prompt": [{"role": "user", "content": "q0"}], "image": object()}, + {"prompt": [{"role": "user", "content": "q1"}], "image": object()}, + {"prompt": [{"role": "user", "content": "q2"}], "image": object()}, + ] + + inputs = trainer._prepare_inputs(generation_batch) + + # The whole eval batch is collated once (no per-accumulation-step slicing) into a tensor dict. + assert len(collated_per_call) == 1 + assert collated_per_call[0] == generation_batch + assert inputs["input_ids"].shape[0] == len(generation_batch) + # Off-policy only: no generation occurred, so no on-policy buffer/log state was created. + assert not hasattr(trainer, "_buffered_inputs") or trainer._buffered_inputs is None + + +def test_eval_vlm_attaches_raw_images_for_teacher_processor(monkeypatch): + """When a teacher processor is configured (cross-arch / ULD), eval must attach raw images and prompts.""" + trainer = GOLDTrainer.__new__(GOLDTrainer) + trainer.accelerator = SimpleNamespace(device=torch.device("cpu"), is_main_process=True) + trainer.model = SimpleNamespace(training=False) + trainer.use_uld_loss = False + trainer.teacher_tokenizer = None + trainer._teacher_processor = object() + + def stub_collator(examples): + return {"input_ids": torch.zeros(len(examples), 1, dtype=torch.long)} + + trainer._vlm_collator = stub_collator + # The exact multimodal-message shape is irrelevant here; pass the prompt through unchanged. + monkeypatch.setattr(gold_trainer_module, "prepare_multimodal_messages", lambda prompt, images: prompt) + + img0, img1 = object(), object() + generation_batch = [ + {"prompt": [{"role": "user", "content": "q0"}], "image": img0}, + {"prompt": [{"role": "user", "content": "q1"}], "images": [img1]}, + ] + + inputs = trainer._prepare_inputs(generation_batch) + + assert inputs["_raw_images"] == [[img0], [img1]] + assert inputs["_raw_prompts"] == [ex["prompt"] for ex in generation_batch] + + def test_on_policy_vlm_without_vllm_collates_only_consumed_slice(monkeypatch): trainer = GOLDTrainer.__new__(GOLDTrainer) trainer.accelerator = SimpleNamespace(device=torch.device("cpu"), is_main_process=True) diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index e8ec71dbc7a..b4fae4c5f55 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -1196,6 +1196,26 @@ def get_train_dataloader(self): @profiling_decorator def _prepare_inputs(self, generation_batch: dict[str, torch.Tensor | Any]) -> dict[str, torch.Tensor | Any]: if not self.model.training: + # Evaluation is off-policy only + if self._vlm_collator is not None: + # Mirror the off-policy slice construction in _fill_buffer: extract raw images and prompts BEFORE + # collation + pending_slice = {"_gold_vlm_lazy_examples": list(generation_batch)} + if self._teacher_processor is not None: + raw_images = [ + ex.get("images") or ([ex["image"]] if "image" in ex else None) for ex in generation_batch + ] + raw_prompts = [ + ( + prepare_multimodal_messages(ex["prompt"], images=imgs) + if imgs is not None + else ex.get("prompt") + ) + for ex, imgs in zip(generation_batch, raw_images, strict=True) + ] + pending_slice["_gold_vlm_raw_images"] = raw_images + pending_slice["_gold_vlm_raw_prompts"] = raw_prompts + return self._materialize_vlm_slice(pending_slice) return generation_batch buffer_steps = self.args.gradient_accumulation_steps @@ -2305,7 +2325,9 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N } ] teacher_prompt_completion_text = self._teacher_processor.apply_chat_template( - raw_prompt + teacher_completion, tokenize=False, add_generation_prompt=False + raw_prompt + teacher_completion, + tokenize=False, + add_generation_prompt=False, ) if teacher_prompt_completion_text.startswith(teacher_prompt_text): teacher_completion_texts.append(teacher_prompt_completion_text[len(teacher_prompt_text) :]) @@ -2326,7 +2348,6 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N teacher_sequences = [] teacher_attention_masks = [] teacher_labels_list = [] - teacher_prompt_token_lengths = [] teacher_sequence_kwargs = defaultdict(list) teacher_sequence_keys = ("token_type_ids", "mm_token_type_ids") @@ -2336,7 +2357,6 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N if eos_token_id is not None and prompt_ids and prompt_ids[-1] == eos_token_id: prompt_ids = prompt_ids[:-1] - teacher_prompt_token_lengths.append(len(prompt_ids)) sequence = list(prompt_ids) sequence.extend(completion_ids) if eos_token_id is not None and eos_token_id not in completion_ids: @@ -2371,7 +2391,6 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N ) teacher_attention_mask = pad(teacher_attention_masks, padding_side="right", padding_value=0).bool() teacher_labels = pad(teacher_labels_list, padding_side="right", padding_value=-100) - teacher_prompt_length = min(teacher_prompt_token_lengths) # Override teacher_forward_kwargs with multimodal keys from teacher processing. teacher_forward_kwargs = { @@ -2389,7 +2408,7 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N teacher_input_ids, teacher_labels, teacher_attention_mask, - teacher_prompt_length, + _, ) = build_teacher_inputs_from_texts( self.teacher_tokenizer, prompt_texts, @@ -2414,17 +2433,6 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N attention_mask=teacher_attention_mask, **teacher_forward_kwargs, ) - - # These are not used for ULD loss but are needed if JSD loss were to be used in this branch - # For VLMs, prompts are left-padded but input_ids are flushed left, so prompts.shape[1] - # would overcount. Derive prompt length from the flushed labels instead. - if self._is_vlm: - student_prompt_length = self._get_min_completion_start_from_labels(inputs["labels"]) - else: - student_prompt_length = inputs["prompts"].shape[1] - shifted_student_logits = outputs_student.logits[:, student_prompt_length - 1 : -1, :] - shifted_teacher_logits = outputs_teacher.logits[:, teacher_prompt_length - 1 : -1, :] - shifted_labels = inputs["labels"][:, student_prompt_length:] else: if self.use_liger_gkd_loss: # Forward only through the base models (avoid lm_head to save memory) From 25d9cadf2a053af232da2b92a58225fef7b04335 Mon Sep 17 00:00:00 2001 From: Strongich Date: Sat, 23 May 2026 12:14:01 +0200 Subject: [PATCH 35/39] fix forward kwargs for cross-tokenizer text-only path --- trl/experimental/gold/gold_trainer.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index 0a496839cfe..890af64a752 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -2404,6 +2404,8 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N self.accelerator.device ) else: + # Text-only cross-tokenizer ULD: teacher inputs are rebuilt independently + teacher_forward_kwargs = {} ( teacher_input_ids, teacher_labels, From 4cfeb3754b7a03ef7c5bc0600580e891cff76e03 Mon Sep 17 00:00:00 2001 From: Strongich Date: Sat, 23 May 2026 12:40:12 +0200 Subject: [PATCH 36/39] fix VLM collator tokenizer & align empty-image handling across rollout paths --- trl/experimental/gold/gold_trainer.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index 890af64a752..60df257227e 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -918,7 +918,7 @@ def __init__( ) data_collator = identity else: - data_collator = DataCollatorForChatML(tokenizer=processing_class, max_length=args.max_length) + data_collator = DataCollatorForChatML(tokenizer=tokenizer, max_length=args.max_length) # Liger fused GKD loss (JSD) self.use_liger_gkd_loss = False @@ -1213,7 +1213,8 @@ def _prepare_inputs(self, generation_batch: dict[str, torch.Tensor | Any]) -> di ) for ex, imgs in zip(generation_batch, raw_images, strict=True) ] - pending_slice["_gold_vlm_raw_images"] = raw_images + has_images = any(imgs is not None for imgs in raw_images) + pending_slice["_gold_vlm_raw_images"] = raw_images if has_images else None pending_slice["_gold_vlm_raw_prompts"] = raw_prompts return self._materialize_vlm_slice(pending_slice) return generation_batch @@ -1475,7 +1476,8 @@ def _fill_buffer( ) for ex, imgs in zip(raw_slices[i], raw_images, strict=True) ] - slice_inputs["_gold_vlm_raw_images"] = raw_images + has_images = any(imgs is not None for imgs in raw_images) + slice_inputs["_gold_vlm_raw_images"] = raw_images if has_images else None slice_inputs["_gold_vlm_raw_prompts"] = raw_prompts else: slice_inputs = slices[i] From 6dfd8fe53f959fe30ac0951a439b8309d3924c55 Mon Sep 17 00:00:00 2001 From: Strongich Date: Sat, 23 May 2026 13:16:15 +0200 Subject: [PATCH 37/39] support VLM eval via GRPO-style prediction_step --- tests/experimental/test_gold_trainer.py | 41 +++++++++++++++++++++---- trl/experimental/gold/gold_trainer.py | 21 +++++++++++-- 2 files changed, 53 insertions(+), 9 deletions(-) diff --git a/tests/experimental/test_gold_trainer.py b/tests/experimental/test_gold_trainer.py index 1cf6a54403f..3d7ca92c6dd 100644 --- a/tests/experimental/test_gold_trainer.py +++ b/tests/experimental/test_gold_trainer.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import contextlib from types import SimpleNamespace import pytest @@ -2081,7 +2082,7 @@ def test_on_policy_vlm_vllm_does_not_duplicate_repeated_sampler_batch(monkeypatc class StubProcessor: @staticmethod - def apply_chat_template(conversation, add_generation_prompt, tokenize, return_dict, padding): + def apply_chat_template(conversation, add_generation_prompt, tokenize, return_dict, processor_kwargs): return { "input_ids": [[1, 2, 3] for _ in conversation], "attention_mask": [[1, 1, 1] for _ in conversation], @@ -2217,11 +2218,10 @@ def stub_collator(examples): def test_eval_vlm_collates_raw_batch_off_policy(): - """During eval the identity collator yields raw dicts; `_prepare_inputs` must collate them off-policy. + """VLM eval (identity collator yields raw dicts) must collate off-policy in `_prepare_inputs`. - Regression test: previously the eval branch returned the raw `list[dict]` unchanged, which crashed `compute_loss` - on `inputs["input_ids"]`. Eval must run the VLM collator over the whole batch (no slicing, no buffering, no - generation). + Regression test for the eval crash: the inherited path indexed the raw `list[dict]`. Eval must run the VLM collator + over the whole batch (no slicing, no buffering, no generation). """ trainer = GOLDTrainer.__new__(GOLDTrainer) trainer.accelerator = SimpleNamespace(device=torch.device("cpu"), is_main_process=True) @@ -2249,7 +2249,7 @@ def stub_collator(examples): assert len(collated_per_call) == 1 assert collated_per_call[0] == generation_batch assert inputs["input_ids"].shape[0] == len(generation_batch) - # Off-policy only: no generation occurred, so no on-policy buffer/log state was created. + # Off-policy only: no generation occurred, so no on-policy buffer state was created. assert not hasattr(trainer, "_buffered_inputs") or trainer._buffered_inputs is None @@ -2281,6 +2281,35 @@ def stub_collator(examples): assert inputs["_raw_prompts"] == [ex["prompt"] for ex in generation_batch] +def test_prediction_step_collates_and_computes_loss(): + """prediction_step must collate via _prepare_inputs and force compute_loss (no raw-input indexing).""" + trainer = GOLDTrainer.__new__(GOLDTrainer) + trainer.model = SimpleNamespace(training=False) + trainer.compute_loss_context_manager = contextlib.nullcontext + + seen = {} + + def fake_prepare_inputs(batch): + seen["prepared"] = batch + return {"input_ids": torch.zeros(len(batch), 1, dtype=torch.long)} + + def fake_compute_loss(model, inputs): + seen["compute_inputs"] = inputs + return torch.tensor([2.0, 4.0]) + + trainer._prepare_inputs = fake_prepare_inputs + trainer.compute_loss = fake_compute_loss + + raw_batch = [{"prompt": "q0"}, {"prompt": "q1"}] + loss, logits, labels = trainer.prediction_step(object(), raw_batch, prediction_loss_only=True) + + # Raw list is routed through _prepare_inputs before any dict indexing, then compute_loss runs. + assert seen["prepared"] is raw_batch + assert "input_ids" in seen["compute_inputs"] + assert loss.item() == 3.0 + assert logits is None and labels is None + + def test_on_policy_vlm_without_vllm_collates_only_consumed_slice(monkeypatch): trainer = GOLDTrainer.__new__(GOLDTrainer) trainer.accelerator = SimpleNamespace(device=torch.device("cpu"), is_main_process=True) diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index 60df257227e..c17f18cb159 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -1196,10 +1196,13 @@ def get_train_dataloader(self): @profiling_decorator def _prepare_inputs(self, generation_batch: dict[str, torch.Tensor | Any]) -> dict[str, torch.Tensor | Any]: if not self.model.training: - # Evaluation is off-policy only + # Evaluation is off-policy (no generation): the student never samples, both models are forwarded over + # the dataset's ground-truth prompt+completion and the distillation loss is taken over the completion. + # For text the collated tensor dict is consumed directly. For VLMs the identity collator yields raw + # dicts (to preserve PIL images for the train-time on-policy path), so collate them here -- mirroring + # the off-policy slice construction in _fill_buffer -- including the raw images/prompts the cross-arch + # / ULD teacher processor needs. if self._vlm_collator is not None: - # Mirror the off-policy slice construction in _fill_buffer: extract raw images and prompts BEFORE - # collation pending_slice = {"_gold_vlm_lazy_examples": list(generation_batch)} if self._teacher_processor is not None: raw_images = [ @@ -2670,6 +2673,18 @@ def _get_liger_zero3_lm_head_gather_ctx(self, model: nn.Module): params.append(teacher_head.bias) return deepspeed.zero.GatheredParameters(params, modifier_rank=None) + # During eval, Trainer calls prediction_step. The inherited SFT prediction_step indexes the raw inputs before + # collation, which breaks the VLM identity-collator path (inputs is a list of raw examples). We override it to + # collate via _prepare_inputs and force compute_loss, evaluating the off-policy distillation loss over the + # ground-truth completion (no generation). + def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys: list[str] | None = None): + inputs = self._prepare_inputs(inputs) + with torch.no_grad(): + with self.compute_loss_context_manager(): + loss = self.compute_loss(model, inputs) + loss = loss.mean().detach() + return loss, None, None + @profiling_decorator def training_step( self, From 325f897b672de41ee40ebeb7af7f7e37111dd7b3 Mon Sep 17 00:00:00 2001 From: Strongich Date: Sat, 23 May 2026 13:58:29 +0200 Subject: [PATCH 38/39] add eval to GOLD VLM examples --- examples/scripts/gold_vlm.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/examples/scripts/gold_vlm.py b/examples/scripts/gold_vlm.py index c46e95c01ef..8ea030ef474 100644 --- a/examples/scripts/gold_vlm.py +++ b/examples/scripts/gold_vlm.py @@ -132,6 +132,11 @@ def convert_to_rgb(example): dataset = dataset.map(convert_to_rgb) dataset = dataset.map(make_conversation) + # Hold out 5% for evaluation + dataset = dataset.train_test_split(test_size=0.05, seed=42) + train_dataset = dataset["train"] + eval_dataset = dataset["test"] + # ────────────────────────────────────────────── # Training config # ────────────────────────────────────────────── @@ -162,6 +167,10 @@ def convert_to_rgb(example): max_steps=300, learning_rate=1e-4, warmup_steps=10, + # Evaluation + per_device_eval_batch_size=2, + eval_strategy="steps", + eval_steps=50, # Precision bf16=True, # Logging @@ -177,7 +186,8 @@ def convert_to_rgb(example): model=student_model, teacher_model=teacher_model, args=args, - train_dataset=dataset, + train_dataset=train_dataset, + eval_dataset=eval_dataset, processing_class=processor, peft_config=peft_config, ) From c187966ac001571f8f781b28bdfd8b431cf3177b Mon Sep 17 00:00:00 2001 From: Strongich Date: Tue, 26 May 2026 12:11:59 +0200 Subject: [PATCH 39/39] drop defensive guards & dedup VLM setup --- tests/experimental/test_gold_trainer.py | 2 +- trl/experimental/gold/gold_trainer.py | 177 ++++++++++-------------- 2 files changed, 77 insertions(+), 102 deletions(-) diff --git a/tests/experimental/test_gold_trainer.py b/tests/experimental/test_gold_trainer.py index 3d7ca92c6dd..2cc482a3044 100644 --- a/tests/experimental/test_gold_trainer.py +++ b/tests/experimental/test_gold_trainer.py @@ -2272,7 +2272,7 @@ def stub_collator(examples): img0, img1 = object(), object() generation_batch = [ {"prompt": [{"role": "user", "content": "q0"}], "image": img0}, - {"prompt": [{"role": "user", "content": "q1"}], "images": [img1]}, + {"prompt": [{"role": "user", "content": "q1"}], "image": img1}, ] inputs = trainer._prepare_inputs(generation_batch) diff --git a/trl/experimental/gold/gold_trainer.py b/trl/experimental/gold/gold_trainer.py index c17f18cb159..a0e3da88895 100644 --- a/trl/experimental/gold/gold_trainer.py +++ b/trl/experimental/gold/gold_trainer.py @@ -849,17 +849,28 @@ def __init__( # VLMs so that both receive images and multimodal inputs. self._teacher_processor = None self._is_cross_architecture_vlm = False - if self._is_vlm and isinstance(teacher_model, str): - # Teacher not yet instantiated -- validate it's a VLM - teacher_proc = AutoProcessor.from_pretrained(teacher_model) - if not isinstance(teacher_proc, ProcessorMixin): - raise ValueError( - "VLM distillation requires both student and teacher to be vision-language models. " - "The student has a `ProcessorMixin` but the teacher does not." - ) + if self._is_vlm: + if isinstance(teacher_model, str): + # Teacher not yet instantiated -- validate it's a VLM + teacher_proc = AutoProcessor.from_pretrained(teacher_model) + if not isinstance(teacher_proc, ProcessorMixin): + raise ValueError( + "VLM distillation requires both student and teacher to be vision-language models. " + "The student has a `ProcessorMixin` but the teacher does not." + ) + teacher_model_type = AutoConfig.from_pretrained(teacher_model).model_type + else: + # Teacher already instantiated — check if it looks like a VLM by checking for a vision config + if not hasattr(teacher_model.config, "vision_config"): + raise ValueError( + "VLM distillation requires both student and teacher to be vision-language models. " + "The student has a `ProcessorMixin` but the teacher model does not appear to be a VLM " + "(missing `vision_config`)." + ) + teacher_model_type = teacher_model.config.model_type + # Check for cross-architecture VLM distillation student_model_type = model.config.model_type if not isinstance(model, str) else None - teacher_model_type = AutoConfig.from_pretrained(teacher_model).model_type is_cross_architecture = student_model_type and teacher_model_type != student_model_type self._is_cross_architecture_vlm = is_cross_architecture if is_cross_architecture: @@ -869,28 +880,11 @@ def __init__( "model's processor, which may increase memory usage and computation time." ) if is_cross_architecture or args.use_uld_loss: - self._teacher_processor = teacher_proc - elif self._is_vlm and not isinstance(teacher_model, str): - # Teacher already instantiated — check if it looks like a VLM by checking for a vision config - if not hasattr(teacher_model, "config") or not hasattr(teacher_model.config, "vision_config"): - raise ValueError( - "VLM distillation requires both student and teacher to be vision-language models. " - "The student has a `ProcessorMixin` but the teacher model does not appear to be a VLM " - "(missing `vision_config`)." + self._teacher_processor = ( + teacher_proc + if isinstance(teacher_model, str) + else AutoProcessor.from_pretrained(teacher_model.config._name_or_path) ) - # Check for cross-architecture VLM distillation - student_model_type = model.config.model_type if not isinstance(model, str) else None - teacher_model_type = teacher_model.config.model_type - is_cross_architecture = student_model_type and teacher_model_type != student_model_type - self._is_cross_architecture_vlm = is_cross_architecture - if is_cross_architecture: - warnings.warn( - f"Cross-architecture VLM distillation detected: student is '{student_model_type}', " - f"teacher is '{teacher_model_type}'. Images will be processed separately through each " - "model's processor, which may increase memory usage and computation time." - ) - if is_cross_architecture or args.use_uld_loss: - self._teacher_processor = AutoProcessor.from_pretrained(teacher_model.config._name_or_path) if self._is_cross_architecture_vlm and not args.use_uld_loss: raise ValueError( "Cross-architecture VLM distillation (student and teacher have different `model_type`) is not " @@ -976,7 +970,7 @@ def __init__( self.teacher_tokenizer.pad_token = self.teacher_tokenizer.eos_token elif args.use_uld_loss and args.teacher_tokenizer_name_or_path is not None: self.teacher_tokenizer = AutoTokenizer.from_pretrained(args.teacher_tokenizer_name_or_path) - if not hasattr(self.teacher_tokenizer, "pad_token") or self.teacher_tokenizer.pad_token is None: + if self.teacher_tokenizer.pad_token is None: self.teacher_tokenizer.pad_token = self.teacher_tokenizer.eos_token # Hybrid ULD loss configuration is handled in ULDLoss class @@ -1193,6 +1187,30 @@ def get_train_dataloader(self): return self.accelerator.prepare(DataLoader(train_dataset, **dataloader_params)) + def _extract_images_and_prompts(self, examples: list[dict]) -> tuple[list | None, list]: + """ + Extract per-example images and build prompts with multimodal messages, mirroring GRPOTrainer. + + Returns `(images, prompts)` where `images` is a per-example list (entries may be `None`), or `None` when the + batch carries no images, and `prompts` are the prepared multimodal messages. + """ + if "images" in examples[0]: + images = [example.get("images") for example in examples] + elif "image" in examples[0]: + images = [[example.get("image")] if example.get("image") is not None else None for example in examples] + else: + images = None + if images is not None and all(img_list is None or img_list == [] for img_list in images): + images = None + + prompts = [example["prompt"] for example in examples] + if images is not None: + prompts = [ + prepare_multimodal_messages(prompt, images=img_list) + for prompt, img_list in zip(prompts, images, strict=True) + ] + return images, prompts + @profiling_decorator def _prepare_inputs(self, generation_batch: dict[str, torch.Tensor | Any]) -> dict[str, torch.Tensor | Any]: if not self.model.training: @@ -1205,19 +1223,8 @@ def _prepare_inputs(self, generation_batch: dict[str, torch.Tensor | Any]) -> di if self._vlm_collator is not None: pending_slice = {"_gold_vlm_lazy_examples": list(generation_batch)} if self._teacher_processor is not None: - raw_images = [ - ex.get("images") or ([ex["image"]] if "image" in ex else None) for ex in generation_batch - ] - raw_prompts = [ - ( - prepare_multimodal_messages(ex["prompt"], images=imgs) - if imgs is not None - else ex.get("prompt") - ) - for ex, imgs in zip(generation_batch, raw_images, strict=True) - ] - has_images = any(imgs is not None for imgs in raw_images) - pending_slice["_gold_vlm_raw_images"] = raw_images if has_images else None + raw_images, raw_prompts = self._extract_images_and_prompts(list(generation_batch)) + pending_slice["_gold_vlm_raw_images"] = raw_images pending_slice["_gold_vlm_raw_prompts"] = raw_prompts return self._materialize_vlm_slice(pending_slice) return generation_batch @@ -1468,19 +1475,8 @@ def _fill_buffer( # mutates examples in place (pops "image", overwrites "prompt"). slice_inputs = {"_gold_vlm_lazy_examples": raw_slices[i]} if self._teacher_processor is not None: - raw_images = [ - ex.get("images") or ([ex["image"]] if "image" in ex else None) for ex in raw_slices[i] - ] - raw_prompts = [ - ( - prepare_multimodal_messages(ex["prompt"], images=imgs) - if imgs is not None - else ex.get("prompt") - ) - for ex, imgs in zip(raw_slices[i], raw_images, strict=True) - ] - has_images = any(imgs is not None for imgs in raw_images) - slice_inputs["_gold_vlm_raw_images"] = raw_images if has_images else None + raw_images, raw_prompts = self._extract_images_and_prompts(raw_slices[i]) + slice_inputs["_gold_vlm_raw_images"] = raw_images slice_inputs["_gold_vlm_raw_prompts"] = raw_prompts else: slice_inputs = slices[i] @@ -1615,25 +1611,8 @@ def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_in for slice_idx in on_policy_indices: raw_examples = raw_slices[slice_idx] - # Extract raw PIL images from examples (like GRPOTrainer) - if "images" in raw_examples[0]: - images = [example.get("images") for example in raw_examples] - elif "image" in raw_examples[0]: - images = [ - [example.get("image")] if example.get("image") is not None else None for example in raw_examples - ] - else: - images = None - if images is not None and all(img_list is None or img_list == [] for img_list in images): - images = None - - # Extract prompts and prepare multimodal messages - prompts = [example["prompt"] for example in raw_examples] - if images is not None: - prompts = [ - prepare_multimodal_messages(prompt, images=img_list) - for prompt, img_list in zip(prompts, images, strict=True) - ] + # Extract raw PIL images and build prompts with multimodal messages (like GRPOTrainer) + images, prompts = self._extract_images_and_prompts(raw_examples) # Normalize string content to content blocks for VLM processors that don't handle plain strings # copied from GRPOTrainer @@ -2288,6 +2267,9 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N # the teacher processor below. teacher_forward_kwargs = student_forward_kwargs + teacher_labels = None + teacher_input_ids = None + if self.use_uld_loss and self.teacher_tokenizer is not None: if "original_prompt_text" in inputs and "original_completion_text" in inputs: prompt_texts = inputs["original_prompt_text"] @@ -2542,18 +2524,14 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N if self.use_uld_loss: student_input_ids = inputs["input_ids"] - teacher_labels_for_loss = teacher_labels if "teacher_labels" in locals() else inputs["labels"] - teacher_input_ids_for_loss = teacher_input_ids if "teacher_input_ids" in locals() else inputs["input_ids"] + teacher_labels_for_loss = teacher_labels if teacher_labels is not None else inputs["labels"] + teacher_input_ids_for_loss = teacher_input_ids if teacher_input_ids is not None else inputs["input_ids"] student_labels = inputs["labels"].clone() if self.pad_token_id is not None: student_labels[student_labels == self.pad_token_id] = -100 - if ( - hasattr(self, "teacher_tokenizer") - and hasattr(self.teacher_tokenizer, "pad_token_id") - and self.teacher_tokenizer.pad_token_id is not None - ): + if self.teacher_tokenizer is not None and self.teacher_tokenizer.pad_token_id is not None: teacher_labels[teacher_labels == self.teacher_tokenizer.pad_token_id] = -100 loss = self.uld_loss_fn( @@ -2565,24 +2543,21 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N teacher_input_ids=teacher_input_ids_for_loss, ) - if hasattr(self.uld_loss_fn, "last_matched_loss") and hasattr(self.uld_loss_fn, "last_unmatched_loss"): - ga = max(1, int(self.args.gradient_accumulation_steps)) - step_eq = 1.0 / ga - matched_val = ( - self.uld_loss_fn.last_matched_loss.item() - if self.uld_loss_fn.last_matched_loss is not None - else 0.0 - ) - unmatched_val = ( - self.uld_loss_fn.last_unmatched_loss.item() - if self.uld_loss_fn.last_unmatched_loss is not None - else 0.0 - ) + ga = max(1, int(self.args.gradient_accumulation_steps)) + step_eq = 1.0 / ga + matched_val = ( + self.uld_loss_fn.last_matched_loss.item() if self.uld_loss_fn.last_matched_loss is not None else 0.0 + ) + unmatched_val = ( + self.uld_loss_fn.last_unmatched_loss.item() + if self.uld_loss_fn.last_unmatched_loss is not None + else 0.0 + ) - self._matched_sum += matched_val - self._unmatched_sum += unmatched_val - self._matched_step_eq += step_eq - self._unmatched_step_eq += step_eq + self._matched_sum += matched_val + self._unmatched_sum += unmatched_val + self._matched_step_eq += step_eq + self._unmatched_step_eq += step_eq empty_cache() @@ -2739,7 +2714,7 @@ def log(self, logs: dict[str, float], start_time: float | None = None) -> None: metrics = {key: sum(val) / len(val) for key, val in self._metrics[mode].items()} # average the metrics if mode == "train": - device = self.accelerator.device if hasattr(self.accelerator, "device") else torch.device("cpu") + device = self.accelerator.device vec = torch.tensor( [ self._on_policy_loss_total,