[GOLD] VLM support for GOLDTrainer

[Strongich]

Clear PR, based on #5461 with changes introduced in 2ac060e

Adds VLM support to GOLDTrainer:

• Supports same-family VLM distillation with JSD loss when student and teacher share compatible tokenization/image-token semantics.
• Supports cross-family VLM distillation with ULD loss by processing images separately through the student and teacher processors.
• Preserves raw PIL images through the dataloader with an identity collator, then materializes VLM batches only for the current accumulation slice.
• Adds VLM on-policy generation paths for both local generation and vLLM-backed generation where the vLLM backend supports the selected model.
• Adds teacher-side VLM input construction for ULD, including a fix for a silent alignment bug: teacher completions are now rendered with the teacher chat template instead of approximating them from generated text plus a manually appended EOS.
• Adds examples/scripts/gold_vlm.py with documented same-family JSD and cross-family ULD examples.
• Adds tests covering VLM collation, label masking, raw image preservation, cross-architecture validation, teacher processor setup, vLLM behavior, lazy slice materialization, and ULD alignment regressions.

Motivation

The GOLD algorithm has no theoretical constraints against VLM-to-VLM distillation -- the barriers were purely engineering (incompatible image token formats, different tokenizers, raw image handling through the dataloader).

Key changes

• GOLDTrainer detects VLM datasets and uses an identity collator to preserve raw PIL images through the dataloader
• For cross-architecture pairs, a _teacher_processor is stored and used in compute_loss to build teacher-compatible vision tensors from raw images
• Auto-resolves teacher_tokenizer_name_or_path
• Added examples/scripts/gold_vlm.py with two documented usage examples (same-family JSD + vLLM, cross-family ULD)
• Added tests for VLM collator (label masking, completion preservation), cross-architecture detection (rejects JSD, stores teacher processor for different archs, skips it for same arch), VLM + vLLM init (copied from the LLM example), rejects LLM teacher with vision dataset
• VLM handling (identity collator, raw image storage, vLLM multimodal path) is borrowed (where it was possible) from SFTTrainer and GRPOTrainer

Note

• I didn't add VLM usage examples to docs/source/gold_trainer.md -- will add if that's desirable, just let me know.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline, Pull Request section?
• Was this discussed/approved via a GitHub issue? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

AI writing disclosure

We welcome the use of AI tools to help with contributions. For transparency and to help us improve our review process, please indicate the level of AI involvement in this PR.

• No AI usage: the PR was written entirely by a human.
• AI-assisted: some parts were suggested or improved by AI, but the PR was written and reviewed by a human.
• AI-generated: the PR was mostly or fully generated by an AI tool.

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag members/contributors who may be interested in your PR.

@qgallouedec @kashif

---

Note

Medium Risk
Large changes to the experimental GOLD training loop (buffering, generation, eval, loss paths) with multimodal and cross-architecture edge cases; mitigated by extensive tests but still affects core distillation behavior.

Overview
GOLDTrainer now supports vision-language model distillation end-to-end: same-family pairs use standard JSD (and optional Liger) with shared pixel_values; different model_type pairs require use_uld_loss=True and route images through separate student/teacher processors.

Vision batches keep raw PIL images in the dataloader via an identity collator, then DataCollatorForVisionLanguageChatML materializes tensors per gradient-accumulation slice (off-policy, local on-policy, and vLLM on-policy). Cross-tokenizer ULD gets _build_teacher_vlm_inputs (teacher chat template + pixel_values) and _raw_images / _raw_prompts on the batch. Init-time guards reject LLM+vision data, non-VLM teachers, keep_end truncation for VLMs, custom collators, and Liger+ULD together.

Adds trl/experimental/gold/gold_vlm.py (GEOQA example), a VLM section in gold_trainer.md, and a large test_gold_trainer.py suite including slow smoke steps for JSD/ULD/Liger on tiny VLMs.

^{Reviewed by Cursor Bugbot for commit 2a3e1c8. Bugbot is set up for automated code reviews on this repo. Configure here.}

[Strongich]

New training runs from examples/scripts/gold_vlm.py:

1. Qwen 8B -> Qwen 2B, JSD loss, vLLM:

2. Qwen 8B -> Qwen 2B, ULD loss, vLLM: 3. Qwen 8B -> LFM 1.6B, ULD loss, no-vllm:

Results are consistent with #5461 (comment) and kashif#6 (comment)

[kashif]

nice! the scripts for experimental belong in the experimental trainer's folder for now

[Strongich]

To avoid opening a separate discussion or issue, I think I can address this idea here (but if you'd prefer otherwise, I'll create one):

liger_kernel is not yet supported with VLMs. I plan to work on it in 2 stages:

1. add VLM support to GKDTrainer
2. add liger kernel VLM support to both GOLDTrainer and GKDTrainer (in a single PR)

Since GKD functionality is fundamentally a case of GOLD, when use_uld_loss=False and we use same-family models, adding duplicative VLM support to the GKDTrainer is excessive imho.

I think we could later combine everything under a single GOLDTrainer (basically dropping a separate GKDTrainer) to keep the whole distillation logic in one place and avoid having to change the JSD path in two separate places. This way, I wouldn't need to add separate VLM support to GKDTrainer and then update liger_kernel -- I'd only need to implement the second part.

10.06.2026 Update

Implementation of the liger kernel support for VLMs was easier than I anticipated, so I added it to the GOLDTrainer in e9d9836 and also added tests regarding it.
To check whether it works, I ran the training from the example above (Qwen -> Qwen, JSD, vLLM, use_liger_kernel = True). The results were almost identical to those of the same training with use_liger_kernel=False.

[Copilot]

Pull request overview

This PR extends the experimental GOLDTrainer to support vision-language model (VLM) distillation, including both same-family distillation (JSD) and cross-family distillation (ULD), with identity-collator-based raw image preservation and VLM-aware on-policy generation paths (local + vLLM), plus regression tests and an example script.

Changes:

• Introduces a VLM-specific ChatML collator that tokenizes prompts/completions and processes images on-the-fly while preserving raw prompt/completion text for ULD alignment.
• Updates GOLDTrainer’s buffering, generation, and loss paths to handle VLM batches lazily and to build teacher-compatible VLM inputs for ULD (including teacher-side prompt rendering).
• Adds a VLM distillation example script and expands the GOLDTrainer test suite to cover multimodal collation/alignment and VLM training/generation paths.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 6 comments.

File	Description
trl/experimental/utils.py	Adds DataCollatorForVisionLanguageChatML for VLM prompt+completion collation (incl. byte offsets + prompt tensors).
trl/experimental/gold/gold_trainer.py	Adds VLM detection, identity-collator buffering/materialization, VLM on-policy generation, and teacher-side VLM input building for ULD.
tests/experimental/test_gold_trainer.py	Adds extensive regression coverage for VLM collation, alignment, buffering/lazy materialization, vLLM behavior, and smoke tests.
examples/scripts/gold_vlm.py	Provides runnable same-family (JSD + vLLM) and cross-family (ULD) VLM distillation examples on GEOQA_R1V.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Strongich]

@kashif

Results are consistent with before, but ahead of sharing them I ran into a bug, fully described in #6156 (the PR description is fully AI-generated -- I was too lazy to write it up myself 😄). Since it touches the general Chunked NLL loss in SFTTrainer, I've split it out into its own PR. My suggestion is to merge the PR above first, and then this one, since without it ULD distillation for different families models with no vLLM, and chunked NLL Loss fails.

1. Qwen 8B -> Qwen 2B, JSD, vLLM:

2. Qwen 8B -> Qwen 2B, ULD, vLLM: 3. Qwen 8B -> LFM 1.6B, ULD, no vLLM:

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 4 comments.

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 2a3e1c8. Configure here.}

[Strongich]

@kashif since the results in #5969 (comment), the branch picked up a few more commits. None of them touch the training path exercised by those runs -- they're either cosmetic or guard against edge cases that don't occur in a normal run:

• 69faf72 — doc typo fix (docs/source/gold_trainer.md)
• 124adf3 — torch_dtype → dtype rename in docs/example, no behavior change
• 92110ff — adds a fail-fast ValueError when the generation batch size isn't divisible by gradient_accumulation_steps; only triggers on an invalid config, no-op otherwise
• a4c6f4c — rejects a custom data_collator for VLM training with a ValueError; only triggers if one is passed, no-op otherwise
• 03770c0 — fixes DataCollatorForVisionLanguageChatML dropping the image field when it's None (e.g. a text-only example in a mixed dataset); a no-op when every example actually has an image, as in the runs above
• ba9da96 — replaces silent truncation of over-budget VLM prompts with a ValueError; only triggers when a prompt exceeds max_length - max_completion_length, which none of the runs hit
• 2a3e1c8 — version-gated workaround for a transformers 5.3.x processor bug (already fixed in 5.4.0); no-op outside that version range
• ca413ba — test-only change

So the core used to produce those results is unchanged, and the results reported there still stand