diff --git a/examples/scripts/gold_vlm.py b/examples/scripts/gold_vlm.py
new file mode 100644
index 00000000000..8ea030ef474
--- /dev/null
+++ b/examples/scripts/gold_vlm.py
@@ -0,0 +1,196 @@
+# 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 GEOQA_R1V.
+
+# 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 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 LiquidAI/LFM2.5-VL-1.6B \
+ --teacher_model_name Qwen/Qwen3-VL-8B-Instruct \
+ --use_uld_loss \
+ --no-use_vllm
+"""
+
+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 = "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 GEOQA_R1V 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["problem"]},
+ ],
+ },
+ ],
+ "completion": [
+ {
+ "role": "assistant",
+ "content": [{"type": "text", "text": normalize_solution(example["solution"])}],
+ },
+ ],
+ "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="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_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()
+
+ # ──────────────────────────────────────────────
+ # Models
+ # ──────────────────────────────────────────────
+ 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():
+ if "language_model" not in name:
+ param.requires_grad = False
+
+ processor = AutoProcessor.from_pretrained(cli_args.student_model_name, padding_side="left")
+
+ peft_config = LoraConfig(
+ r=16,
+ lora_alpha=32,
+ lora_dropout=0.05,
+ target_modules=r"^.*language_model.*\.(q_proj|k_proj|v_proj)$",
+ )
+
+ # ──────────────────────────────────────────────
+ # Dataset
+ # ──────────────────────────────────────────────
+ 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)
+
+ # 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
+ # ──────────────────────────────────────────────
+ args = GOLDConfig(
+ 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.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=1024,
+ max_length=2048,
+ # Training schedule
+ per_device_train_batch_size=2,
+ gradient_accumulation_steps=4,
+ 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
+ logging_steps=10,
+ log_completions=True,
+ report_to="wandb",
+ )
+
+ # ──────────────────────────────────────────────
+ # Trainer
+ # ──────────────────────────────────────────────
+ trainer = GOLDTrainer(
+ model=student_model,
+ teacher_model=teacher_model,
+ args=args,
+ train_dataset=train_dataset,
+ eval_dataset=eval_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 f216badf6e3..2cc482a3044 100644
--- a/tests/experimental/test_gold_trainer.py
+++ b/tests/experimental/test_gold_trainer.py
@@ -12,16 +12,25 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+import contextlib
from types import SimpleNamespace
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.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")
@@ -271,6 +280,39 @@ 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_dataset():
+ try:
+ 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}")
+
+
+@pytest.fixture
+def vlm_examples(vlm_dataset):
+ return [dict(row) for row in vlm_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"]
@@ -295,13 +337,19 @@ 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]
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]
@@ -319,8 +367,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]]))
@@ -343,7 +397,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: ""}
@@ -430,7 +489,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 = []
@@ -449,6 +513,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()
@@ -474,8 +539,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"])
@@ -500,6 +571,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(
@@ -516,7 +590,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
@@ -576,12 +657,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(),
)
@@ -709,7 +795,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)
@@ -735,7 +824,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,
)
@@ -783,6 +875,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,
@@ -941,3 +1065,1404 @@ 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,
+ 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 _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)
+ 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_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, 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}"
+ )
+
+
+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,
+ 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,
+ )
+
+
+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,
+ 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"
+
+ class DummyTeacherModel:
+ def __init__(self):
+ self.config = SimpleNamespace(vision_config=True, model_type="dummy_vlm")
+ 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,
+ 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)
+
+
+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.config.get_text_config = lambda: self.config
+ 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,
+ 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="")
+ 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", "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
+ 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)]
+ 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
+ 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)]
+ 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
+
+
+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._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.
+
+ `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_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
+ trainer._is_cross_architecture_vlm = False
+ trainer.use_uld_loss = False
+ trainer.teacher_tokenizer = None
+
+ class StubProcessor:
+ @staticmethod
+ 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],
+ }
+
+ @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),
+ "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
+
+ 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 = [
+ 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,
+ )
+
+ trainer._generate_on_policy_vlm_raw(raw_slices, on_policy_indices)
+
+ 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
+
+ # 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:
+ 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
+
+ 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
+
+
+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._is_cross_architecture_vlm = False
+ 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_eval_vlm_collates_raw_batch_off_policy():
+ """VLM eval (identity collator yields raw dicts) must collate off-policy in `_prepare_inputs`.
+
+ 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)
+ 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 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"}], "image": 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_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)
+ 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
+ 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):
+ 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
+
+ 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),
+ "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],
+ }
+
+ trainer._vlm_collator = stub_collator
+
+ class FakeModel:
+ training = True
+
+ @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))
+
+ 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"}], "completion": "gold0", "image": object()}],
+ [{"prompt": [{"role": "user", "content": "q1"}], "completion": "gold1", "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, 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)
+ 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_config.py b/trl/experimental/gold/gold_config.py
index 93e495e8c1f..1766ac91375 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 581ab869242..a0e3da88895 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 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
@@ -39,9 +39,19 @@
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, 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
@@ -52,10 +62,17 @@
RepeatSampler,
create_model_from_path,
disable_dropout_in_model,
+ get_config_model_id,
+ identity,
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
@@ -213,9 +230,14 @@ 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
+ return (
+ teacher_input_ids,
+ teacher_labels,
+ teacher_attention_mask,
+ teacher_prompt_length,
+ )
class ULDLoss(nn.Module):
@@ -223,7 +245,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
@@ -253,7 +281,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.
@@ -281,7 +315,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
@@ -319,7 +358,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.
@@ -437,7 +482,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
@@ -661,11 +710,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
@@ -742,13 +793,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__(
@@ -759,27 +812,117 @@ 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
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
+ # 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
- # Respect a user-provided data_collator; otherwise, provide a ChatML collator that
+ 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.
+ self._teacher_processor = None
+ self._is_cross_architecture_vlm = False
+ 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
+ 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 = (
+ teacher_proc
+ if isinstance(teacher_model, str)
+ else 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 "
+ "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(
+ "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, pick the right collator based on modality.
+ # 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:
- data_collator = DataCollatorForChatML(tokenizer=processing_class, max_length=args.max_length)
+ if self._is_vision_dataset:
+ self._vlm_collator = DataCollatorForVisionLanguageChatML(
+ processor=processing_class,
+ max_length=args.max_length,
+ )
+ data_collator = identity
+ else:
+ data_collator = DataCollatorForChatML(tokenizer=tokenizer, max_length=args.max_length)
# 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,
@@ -807,6 +950,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."
@@ -819,9 +964,13 @@ 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:
+ 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
@@ -879,7 +1028,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,
)
@@ -890,7 +1039,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
@@ -940,7 +1089,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),
@@ -957,6 +1108,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",
@@ -964,6 +1117,16 @@ 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",
+ "spatial_shapes",
+ "token_type_ids",
+ "mm_token_type_ids",
]
if self._signature_columns is None:
self._signature_columns = required_columns
@@ -1024,9 +1187,46 @@ 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:
+ # 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:
+ pending_slice = {"_gold_vlm_lazy_examples": list(generation_batch)}
+ if self._teacher_processor is not 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
buffer_steps = self.args.gradient_accumulation_steps
@@ -1035,9 +1235,112 @@ 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"]
+ 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()
+ }
+
+ 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 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=sequence_dtype,
+ device=new_input_ids.device,
+ )
+ 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
+ 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, 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."""
+ 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 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")
@@ -1048,8 +1351,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(
@@ -1111,9 +1414,47 @@ 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."""
+ 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], 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)]
@@ -1129,9 +1470,18 @@ 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:
+ # Extract raw images and prompts BEFORE collation, since the collator
+ # 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, 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]
- 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(
@@ -1143,7 +1493,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(
@@ -1180,6 +1533,8 @@ def _generate_on_policy_for_slices(
self.vllm_generation.sync_weights()
self._last_vllm_sync_step = self.state.global_step
+ # 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,
@@ -1210,20 +1565,189 @@ 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
+ (
+ 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
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 self._SEQUENCE_KEYS:
+ if k in updated_slice:
+ 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,
+ 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
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."""
+ # 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]
+
+ # 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
+ prompts = [
+ [
+ (
+ {**msg, "content": [{"type": "text", "text": msg["content"]}]}
+ if isinstance(msg.get("content"), str)
+ else msg
+ )
+ for msg in prompt
+ ]
+ 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(
+ conversation=prompts,
+ add_generation_prompt=True,
+ tokenize=True,
+ return_dict=True,
+ processor_kwargs={"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)
+ ]
+
+ 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)
+
+ if not self.use_vllm:
+ # 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)
+ 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)
+ # 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
+
+ if any(img is not None for img in all_images):
+ generate_images = all_images
+ else:
+ generate_images = None
+ _, completion_ids, _, _ = self.vllm_generation.generate(
+ prompts=all_prompt_ids,
+ images=generate_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=True,
+ clean_up_tokenization_spaces=False,
+ )
+ )
+
+ # 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}
+ slice_prompts = {idx: [] for idx in on_policy_indices}
+ slice_prompts_text = {idx: [] for idx in on_policy_indices}
+
+ for i, slice_idx in enumerate(local_slice_indices):
+ 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])
+
+ 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_for_slice):
+ synthetic = dict(example)
+ # 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)
+
+ has_images = any(img is not None for img in images_for_slice)
+ pending_slice = {
+ "_gold_vlm_lazy_examples": synthetic_examples,
+ }
+ 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,
+ )
+
def _process_completions_to_buffer(
self,
slices: list[dict[str, torch.Tensor | Any]],
@@ -1239,7 +1763,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}
@@ -1286,7 +1810,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)
@@ -1306,7 +1834,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),
]
)
@@ -1350,13 +1882,18 @@ 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,
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
@@ -1370,7 +1907,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,
@@ -1401,7 +1938,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,
)
@@ -1421,7 +1958,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,
)
@@ -1440,7 +1977,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"],
@@ -1517,7 +2056,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
@@ -1550,7 +2091,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),
+ ),
}
)
@@ -1573,7 +2118,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]
@@ -1652,7 +2201,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,
)
@@ -1667,6 +2221,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
@@ -1679,12 +2235,45 @@ def generalized_jsd_loss(
else:
return jsd
+ _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 = 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
+
+ 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"]
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
@@ -1698,16 +2287,122 @@ 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 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_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,
+ padding=True,
+ return_tensors="pt",
+ )
+ teacher_completion_token_ids = self._teacher_processor.tokenizer(
+ teacher_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_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]
+
+ sequence = list(prompt_ids)
+ sequence.extend(completion_ids)
+ 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)
+ 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)
+
+ # Override teacher_forward_kwargs with multimodal keys from teacher processing.
+ teacher_forward_kwargs = {
+ 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(
+ self.accelerator.device
+ )
+ else:
+ # Text-only cross-tokenizer ULD: teacher inputs are rebuilt independently
+ teacher_forward_kwargs = {}
+ (
+ teacher_input_ids,
+ teacher_labels,
+ teacher_attention_mask,
+ _,
+ ) = 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)
@@ -1717,6 +2412,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,
+ **student_forward_kwargs,
)
self.teacher_model.eval()
@@ -1724,13 +2420,8 @@ 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,
+ **teacher_forward_kwargs,
)
-
- # 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]
- 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)
@@ -1739,13 +2430,16 @@ 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(
input_ids=inputs["input_ids"],
attention_mask=inputs["attention_mask"],
use_cache=False,
+ **student_forward_kwargs,
)
self.teacher_model.eval()
@@ -1754,13 +2448,16 @@ 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(
input_ids=inputs["input_ids"],
attention_mask=inputs["attention_mask"],
use_cache=False,
+ **teacher_forward_kwargs,
)
student_hidden = student_outputs.last_hidden_state[:, :-1]
@@ -1773,7 +2470,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)
@@ -1795,6 +2494,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"],
+ **student_forward_kwargs,
)
self.teacher_model.eval()
@@ -1802,9 +2502,14 @@ 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"],
+ **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 = 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, :]
shifted_teacher_logits = outputs_teacher.logits[:, prompt_lengths - 1 : -1, :]
shifted_labels = inputs["labels"][:, prompt_lengths:]
@@ -1819,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 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")
- 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(
@@ -1842,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()
@@ -1867,11 +2565,18 @@ 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 = 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 self._SEQUENCE_KEYS:
+ 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
@@ -1880,7 +2585,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
# model.generate() returns full sequences (prompt + completion), so completions start
# after the full padded prompt width.
@@ -1914,7 +2619,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:
@@ -1937,9 +2648,24 @@ 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, 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.
@@ -1974,6 +2700,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:
@@ -1981,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,
diff --git a/trl/experimental/utils.py b/trl/experimental/utils.py
index 8d86c237953..5d500287f13 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,
@@ -38,8 +40,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():
@@ -250,7 +258,11 @@ def __call__(self, examples: list[dict[str, Any]]) -> dict[str, torch.Tensor]:
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)
+ 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 {
@@ -262,6 +274,204 @@ 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
+
+ # 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):
+ if content:
+ parts.append(content)
+ continue
+ turn_parts: list[str] = []
+ for block in content:
+ if isinstance(block, dict) and block.get("type") == "text":
+ text = block.get("text", "")
+ if text:
+ turn_parts.append(text)
+ elif isinstance(block, str):
+ 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]
+
+ # 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.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"]
+ 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 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
+ 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.
+ # 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
+
+
def truncate_right(
input_ids: torch.Tensor, stop_token_id: int, pad_token_id: int
) -> tuple[torch.Tensor, torch.Tensor]:
@@ -314,7 +524,9 @@ def add_bos_token_if_needed(
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]:
chosen_tokens["input_ids"].append(eos_token_id)
@@ -349,7 +561,10 @@ def first_true_indices(bools: torch.Tensor, dtype=torch.long) -> torch.Tensor:
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.