docs/source/gold_trainer.md
General Online Logit Distillation (GOLD) is an extension of Universal Logit Distillation (ULD) that supports student/teacher pairs with different tokenizers. It aligns the textual spans produced by both tokenizers and merges the associated logits so no completion tokens are dropped. This enables cross-tokenizer knowledge distillation, including mixed model families (for example, LLaMA students with Qwen teachers).
Key capabilities:
uld_use_hybrid_loss is enabled, GOLD compares exact vocabulary matches directly and falls back to the original sorted-probability ULD loss for unmatched tokens. This improves stability for students whose vocabularies only partially overlap with the teacher.experimental.gkd.GKDTrainer], so you can combine sequence-level KD, generalized JSD, and cross-tokenizer distillation in a single training run.[!NOTE] GOLD is currently part of the
trl.experimentalnamespace. APIs may change without notice while the feature is iterated on.
The [GOLDTrainer] subclasses [SFTTrainer] and accepts the same datasets as other TRL trainers (lists of ChatML style
messages). Important configuration flags on [GOLDConfig] include:
use_uld_loss – toggles Universal Logit Distillation. Set this to True for cross-tokenizer setups.teacher_tokenizer_name_or_path – required when use_uld_loss=True; GOLD uses the teacher tokenizer to align tokens.uld_use_hybrid_loss, uld_hybrid_matched_weight, uld_hybrid_unmatched_weight – enables and weights the hybrid
matched/unmatched loss.beta, lmbda, seq_kd – inherited from [experimental.gkd.GKDConfig], controlling the generalized JSD interpolation and on-policy
sampling ratio.num_generations, generation_batch_size – control buffered rollout generation across gradient accumulation windows.
generation_batch_size is the number of unique prompts per worker per optimizer step.A minimal end-to-end example:
from datasets import load_dataset
from trl.experimental.gold import GOLDConfig, GOLDTrainer
train_dataset = load_dataset(
"HuggingFaceTB/OpenR1-Math-220k-default-verified",
"all",
split="train[:1024]",
)
trainer = GOLDTrainer(
model="meta-llama/Llama-3.2-1B-Instruct",
teacher_model="Qwen/Qwen2.5-0.5B-Instruct",
args=GOLDConfig(output_dir="gold-model", use_uld_loss=True, teacher_tokenizer_name_or_path="Qwen/Qwen2.5-0.5B-Instruct"),
train_dataset=train_dataset,
)
trainer.train()
For quick-start workflows you can rely on string identifiers as shown above—the trainer will load the model and tokenizer for you. Explicitly instantiating AutoModelForCausalLM, AutoTokenizer, or populating GOLDConfig is recommended only for advanced use cases where you need fine-grained control over initialization.
A more explicit setup might look like this when you need to customise model loading, tokenizer settings, or training arguments:
from datasets import load_dataset
from trl.experimental.gold import GOLDConfig, GOLDTrainer
from transformers import AutoModelForCausalLM, AutoTokenizer
student_name = "meta-llama/Llama-3.2-1B-Instruct"
teacher_name = "Qwen/Qwen2.5-0.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(student_name)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained(student_name)
teacher_model = AutoModelForCausalLM.from_pretrained(teacher_name)
train_dataset = load_dataset(
"HuggingFaceTB/Countdown-Task-GOLD",
"verified_Qwen2.5-0.5B-Instruct",
split="train",
)
training_args = GOLDConfig(
output_dir="gold-model",
per_device_train_batch_size=1,
teacher_model_name_or_path=teacher_name,
teacher_tokenizer_name_or_path=teacher_name,
use_uld_loss=True,
uld_use_hybrid_loss=True,
)
trainer = GOLDTrainer(
model=model,
teacher_model=teacher_model,
args=training_args,
processing_class=tokenizer,
train_dataset=train_dataset,
)
trainer.train()
[!NOTE] GOLD buffers one full optimizer-window generation batch (
per_device_train_batch_size * gradient_accumulation_steps) and reuses it across accumulation steps. If the final batch is undersized, GOLD warns and drops that last batch (Dropping last batch due to unexpected batch size). Setdataloader_drop_last=Trueto avoid this warning.
GOLD requires a conversational language modeling dataset, e.g.:
{"messages": [{"role": "user", "content": "What color is the sky?"},
{"role": "assistant", "content": "It is blue."}]}
GOLDTrainer keeps the raw messages so the ChatML collator can construct prompts and completions with the correct
boundaries.
When student and teacher use different tokenizers, the same text may be split differently:
"Hugging Face" → 1 token"Hugging", " Face" → 2 tokensGOLD aligns these sequences and merges the teacher's multi-token probabilities into a single distribution that can be compared with the student's single-token distribution.
For a teacher sequence of tokens [token₀, token₁, ..., tokenₖ] that maps to a single student token, GOLD computes:
P_merged(y) = P(y | context) × P(token₁ | token₀, context) × ... × P(tokenₖ | ..., context)
where:
P(y | context) is the marginal probability distribution over all vocabulary tokens at the first positionP(tokenᵢ | ..., context) are scalar conditional probabilities of the actual tokens that were generatedKey insight: Only the conditional probabilities of the actual continuation tokens are extracted as scalars. The full marginal distribution at the first position is then scaled by multiplying these scalar probabilities.
This ensures:
Given:
P(x₀): ["HF": 0.6, "is": 0.3, "cool": 0.1]
P(x₁ | "HF"): ["HF": 0.05, "is": 0.9, "cool": 0.05]
If tokens 0 and 1 are merged, and the actual sequence was ["HF", "is"]:
P_merged("HF") = 0.6 × 0.9 = 0.54 ✓ (correct joint probability)
P_merged("is") = 0.3 × 0.9 = 0.27
P_merged("cool") = 0.1 × 0.9 = 0.09
The merged distribution is unnormalized (sums to 0.81), but this is intentional and correct for ULD loss computation, which uses sorting and L1 distance.
Use examples/scripts/gold.py to launch GOLD training from the command line. The script supports full training and LoRA via the standard ModelConfig flags.
python examples/scripts/gold.py \
--model_name_or_path meta-llama/Llama-3.2-1B-Instruct \
--teacher_model_name_or_path Qwen/Qwen2-1.5B-Instruct \
--dataset_name trl-lib/chatbot_arena_completions \
--learning_rate 2e-5 \
--per_device_train_batch_size 4 \
--gradient_accumulation_steps 8 \
--output_dir gold-model \
--num_train_epochs 1 \
--push_to_hub
[GOLDTrainer] supports VLM-to-VLM distillation. Both student and teacher must be vision-language models. To train a VLM, provide a dataset with either an image column (single image per sample) or an images column (list of images per sample). For more information on the expected dataset structure, see the Dataset Format — Vision datasets section.
When the student and teacher share the same architecture and tokenizer (e.g. Qwen3-VL-8B to Qwen3-VL-2B), the standard generalized JSD loss applies directly. When they have different model_type (e.g. Qwen3-VL to LFM2.5-VL), set use_uld_loss=True to enable cross-tokenizer alignment via Universal Logit Distillation. Images are processed separately through each model's processor.
from datasets import load_dataset
import torch
from transformers import AutoModelForImageTextToText, AutoProcessor
from trl.experimental.gold import GOLDConfig, GOLDTrainer
student_name = "Qwen/Qwen3-VL-2B-Instruct"
teacher_name = "Qwen/Qwen3-VL-8B-Instruct"
processor = AutoProcessor.from_pretrained(student_name, padding_side="left")
student_model = AutoModelForImageTextToText.from_pretrained(student_name, dtype=torch.bfloat16)
teacher_model = AutoModelForImageTextToText.from_pretrained(teacher_name, dtype=torch.bfloat16)
train_dataset = load_dataset("trl-lib/llava-instruct-mix", split="train")
trainer = GOLDTrainer(
model=student_model,
teacher_model=teacher_model,
args=GOLDConfig(
output_dir="gold-vlm-model",
max_length=None,
teacher_model_name_or_path=teacher_name,
use_uld_loss=False,
),
train_dataset=train_dataset,
processing_class=processor,
)
trainer.train()
For cross-family distillation, set use_uld_loss=True and teacher_tokenizer_name_or_path to the teacher model name.
Use trl/experimental/gold/gold_vlm.py to launch GOLD VLM training from the command line:
# Same-family distillation (JSD loss, vLLM enabled)
accelerate launch trl/experimental/gold/gold_vlm.py \
--student_model_name Qwen/Qwen3-VL-2B-Instruct \
--teacher_model_name Qwen/Qwen3-VL-8B-Instruct
# Cross-family distillation (ULD loss, local generation)
accelerate launch trl/experimental/gold/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
[!TIP] For VLMs,
truncation_mode='keep_end'is not supported because image tokens reside in the prompt portion of the sequence and may be silently dropped. Usetruncation_mode='keep_start'(the default) or setmax_length=Nonein the [GOLDConfig]. This allows the model to process the full sequence length without truncating image tokens.pythonGOLDConfig(max_length=None, ...)Only use
max_lengthwhen you've verified that truncation won't remove image tokens for the entire dataset.
[!NOTE] Cross-architecture VLM distillation requires
use_uld_loss=True. The trainer will raise an error if you attempt cross-architecture distillation without ULD loss.
[[autodoc]] experimental.gold.GOLDTrainer - train - generate_on_policy_outputs - save_model - push_to_hub
[[autodoc]] experimental.gold.GOLDConfig