Treinador de Destilação Online de Logits Geral (GOLD)

Visão geral

General Online Logit Distillation (GOLD) é uma extensão do Universal Logit Distillation (ULD) que suporta pares de student/teacher com tokenizers diferentes. Ele alinha os spans textuais produzidos por ambos os tokenizers e mescla os logits associados para que nenhum token de conclusão seja descartado. Isso permite a destilação de conhecimento entre tokenizers, incluindo famílias de modelos mistas (por exemplo, students LLaMA com teachers Qwen).

Principais recursos:

Alinhamento entre tokenizadores – O GOLD decodifica incrementalmente os tokens do aluno e do professor, agrupa passagens com o mesmo texto visível e mescla as probabilidades dentro de cada grupo. Isso garante que os termos de perda sejam calculados sobre a conclusão completa, mesmo quando as fronteiras dos tokens diferem.
Perda ULD híbrida – quando uld_use_hybrid_loss está habilitado, o GOLD compara diretamente as correspondências exatas do vocabulário e recorre à perda ULD de probabilidade ordenada original para os tokens não correspondidos. Isso melhora a estabilidade para alunos cujos vocabulários se sobrepõem apenas parcialmente ao do professor.
Integração sem emendas com o GKD – o GOLD herda o agendamento on-policy vs. off-policy do experimental.gkd.GKDTrainer, permitindo combinar KD em nível de sequência, JSD generalizado e destilação entre tokenizadores em uma única execução de treinamento.

O GOLD faz parte atualmente do namespace trl.experimental. As APIs podem ser alteradas sem aviso prévio enquanto o recurso estiver em desenvolvimento.

Dicas de uso

O experimental.gold.GOLDTrainer é uma subclasse de SFTTrainer e aceita os mesmos conjuntos de dados que outros treinadores do TRL (listas de mensagens no estilo ChatML). As principais flags de configuração em experimental.gold.GOLDConfig incluem:

use_uld_loss – alterna a Universal Logit Distillation. Defina este valor como True para configurações com tokenizadores diferentes.
teacher_tokenizer_name_or_path – obrigatório quando use_uld_loss=True; o GOLD usa o tokenizer do teacher para alinhar os tokens.
uld_use_hybrid_loss, uld_hybrid_matched_weight, uld_hybrid_unmatched_weight – habilita e pondera a perda híbrida correspondente/não correspondente.
beta, lmbda, seq_kd – herdados de experimental.gkd.GKDConfig, controlando a interpolação do JSD generalizado e a taxa de amostragem on-policy.
num_generations, generation_batch_size – controlam a geração de rollout em buffer entre as janelas de acúmulo de gradiente. generation_batch_size é o número de prompts únicos por worker por passo do otimizador.

Um exemplo mínimo de ponta a ponta:

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()

Para fluxos de trabalho de início rápido, você pode confiar nos identificadores de string, conforme mostrado acima — o trainer carregará o modelo e o tokenizer para você. A instanciação explícita de AutoModelForCausalLM, AutoTokenizer ou a configuração de GOLDConfig é recomendada apenas para casos de uso avançados em que você precisa de controle fino sobre a inicialização.

Um setup mais explícito pode ser assim quando você precisa personalizar o carregamento do modelo, as configurações do tokenizer ou os argumentos de treinamento:

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()

GOLD faz o buffer de um lote de geração completo da janela do otimizador (per_device_train_batch_size * gradient_accumulation_steps) e o reutiliza entre as etapas de acumulação. Se o lote final tiver tamanho insuficiente, o GOLD emite um aviso e descarta esse último lote (Dropping last batch due to unexpected batch size). Defina dataloader_drop_last=True para evitar esse aviso.

Tipo de dataset esperado

O GOLD requer um conjunto de dados de conversação de modelagem de linguagem, por exemplo:

{"messages": [{"role": "user", "content": "What color is the sky?"},
              {"role": "assistant", "content": "It is blue."}]}

GOLDTrainer mantém as mensagens brutas para que o collator do ChatML possa construir prompts e completions com as fronteiras corretas.

Como o Token Merging Funciona

Quando o aluno e o professor usam tokenizadores diferentes, o mesmo texto pode ser dividido de maneiras diferentes:

Aluno: "Hugging Face" → 1 token
Professor: "Hugging", " Face" → 2 tokens

O GOLD alinha essas sequências e mescla as probabilidades de múltiplos tokens do professor em uma única distribuição que pode ser comparada com a distribuição de token único do aluno.

Fusão de Probabilidades

Para uma sequência de tokens do professor [token₀, token₁, ..., tokenₖ] que mapeia para um único token do aluno, o GOLD calcula:

P_merged(y) = P(y | context) × P(token₁ | token₀, context) × ... × P(tokenₖ | ..., context)

onde:

P(y | context) é a distribuição de probabilidade marginal sobre todos os tokens do vocabulário na primeira posição
P(tokenᵢ | ..., context) são probabilidades condicionais escalares dos tokens reais que foram gerados

Insight principal: Somente as probabilidades condicionais dos tokens de continuação reais são extraídas como escalares. A distribuição marginal completa na primeira posição é então escalada multiplicando essas probabilidades escalares.

Isso garante:

Probabilidade conjunta correta para a sequência gerada real (pela regra da cadeia)
Aproximação razoável para tokens contrafatuais (escalados pela mesma probabilidade de continuação)
Distribuições não normalizadas que preservam as probabilidades relativas corretas para o cálculo da perda ULD

Exemplo

Dado:

P(x₀):         ["HF": 0.6,  "is": 0.3,  "cool": 0.1]
P(x₁ | "HF"):  ["HF": 0.05, "is": 0.9,  "cool": 0.05]

Se os tokens 0 e 1 forem mesclados e a sequência real for ["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

A distribuição mesclada não está normalizada (soma 0,81), mas isso é intencional e correto para o cálculo da perda ULD, que usa ordenação e distância L1.

Script de exemplo

Use examples/gold_chatbot_arena/gold_chatbot_arena.py para iniciar o treinamento do GOLD pela linha de comando. O script suporta treinamento completo e LoRA por meio das flags padrão de ModelConfig.

python examples/gold_chatbot_arena/gold_chatbot_arena.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

Treinando Modelos de Visão e Linguagem

experimental.gold.GOLDTrainer suporta destilação de VLM para VLM. Tanto o aluno quanto o professor devem ser modelos de visão-linguagem. Para treinar um VLM, forneça um conjunto de dados com uma coluna image (uma única imagem por amostra) ou uma coluna images (lista de imagens por amostra). Para mais informações sobre a estrutura esperada do conjunto de dados, consulte a seção Formato do conjunto de dados — Conjuntos de dados de visão.

Quando o aluno e o professor compartilham a mesma arquitetura e tokenizer (por exemplo, Qwen3-VL-8B para Qwen3-VL-2B), a perda JSD generalizada padrão se aplica diretamente. Quando eles possuem model_type diferentes (por exemplo, Qwen3-VL para LFM2.5-VL), defina use_uld_loss=True para habilitar o alinhamento entre tokenizers por meio da Universal Logit Distillation. As imagens são processadas separadamente pelo processor de cada modelo.

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()

Para destilação entre famílias, defina use_uld_loss=True e teacher_tokenizer_name_or_path com o nome do modelo teacher.

Use examples/gold_qwen3_vl/gold_qwen3_vl.py para iniciar o treinamento do GOLD VLM pela linha de comando:

# Same-family distillation (JSD loss, vLLM enabled)
torchrun --nproc_per_node 8 examples/gold_qwen3_vl/gold_qwen3_vl.py \
    --student_model_name Qwen/Qwen3-VL-2B-Instruct \
    --teacher_model_name Qwen/Qwen3-VL-8B-Instruct

# Cross-family distillation (ULD loss, local generation)
torchrun --nproc_per_node 8 examples/gold_qwen3_vl/gold_qwen3_vl.py \
    --student_model_name LiquidAI/LFM2.5-VL-1.6B \
    --teacher_model_name Qwen/Qwen3-VL-8B-Instruct \
    --use_uld_loss \
    --no-use_vllm

Para VLMs, truncation_mode='keep_end' não é suportado porque os tokens de imagem estão na parte do prompt da sequência e podem ser descartados silenciosamente. Use truncation_mode='keep_start' (o padrão) ou defina max_length=None no experimental.gold.GOLDConfig. Isso permite que o modelo processe o comprimento total da sequência sem truncar os tokens de imagem.
GOLDConfig(max_length=None, ...)
Use max_length apenas quando você tiver verificado que a truncagem não removerá os tokens de imagem para todo o conjunto de dados.

A destilação de VLM entre arquiteturas diferentes requer use_uld_loss=True. O trainer lançará um erro se você tentar destilação entre arquiteturas diferentes sem a perda ULD.

GOLDTrainer

classe trl.experimental.gold.GOLDTrainer

< source >

( model: typing.Union[transformers.modeling_utils.PreTrainedModel, torch.nn.Module, str, NoneType] = Noneteacher_model: typing.Union[transformers.modeling_utils.PreTrainedModel, torch.nn.Module, str] = Noneargs: trl.experimental.gold.gold_config.GOLDConfig | None = Nonedata_collator: collections.abc.Callable[[list[typing.Any]], dict[str, typing.Any]] | None = Nonetrain_dataset: datasets.arrow_dataset.Dataset | None = Noneeval_dataset: datasets.arrow_dataset.Dataset | dict[str, datasets.arrow_dataset.Dataset] | None = Noneprocessing_class: transformers.tokenization_utils_base.PreTrainedTokenizerBase | transformers.image_processing_utils.BaseImageProcessor | transformers.feature_extraction_utils.FeatureExtractionMixin | transformers.processing_utils.ProcessorMixin | None = Nonecompute_metrics: collections.abc.Callable[[transformers.trainer_utils.EvalPrediction], dict] | None = Nonecallbacks: list[transformers.trainer_callback.TrainerCallback] | None = Noneoptimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None)preprocess_logits_for_metrics: collections.abc.Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = Nonepeft_config: typing.Optional[ForwardRef('PeftConfig')] = None )

treino

< source >

Parâmetros

resume_from_checkpoint (str ou bool, opcional) — Se for um str, caminho local para um checkpoint salvo por uma instância anterior de Trainer. Se for um bool e for igual a True, carrega o último checkpoint em args.output_dir salvo por uma instância anterior de Trainer. Se presente, o treinamento será retomado a partir dos estados do modelo/otimizador/agendador carregados aqui.
trial (optuna.Trial ou dict[str, Any], opcional) — A execução do trial ou o dicionário de hiperparâmetros para a busca de hiperparâmetros.
ignore_keys_for_eval (list[str], opcional) — Uma lista de chaves na saída do seu modelo (se for um dicionário) que devem ser ignoradas ao coletar previsões para avaliação durante o treinamento.

Retorna

~trainer_utils.TrainOutput

Objeto contendo a contagem de etapas globais, a perda de treinamento e as métricas.

Ponto de entrada principal do treinamento.

generate_on_policy_outputs

< source >

( modelinputsgeneration_config )

save_model

< source >

( output_dir: str | None = None_internal_call: bool = False )

Salvará o modelo, para que você possa recarregá-lo usando from_pretrained().

Salvará apenas a partir do processo principal.

push_to_hub

< source >

( commit_message: str | None = 'End of training'blocking: bool = Truetoken: str | None = Nonerevision: str | None = None**kwargs )

Parâmetros

commit_message (str, opcional, padrão "End of training") — Mensagem a ser commitada durante o push.
blocking (bool, opcional, padrão True) — Se a função deve retornar apenas quando o git push tiver terminado.
token (str, opcional, padrão None) — Token com permissão de escrita para sobrescrever os argumentos originais do Trainer.
revision (str, opcional) — A revisão git de onde o commit será feito. O padrão é o topo do ramo “main”.
kwargs (dict[str, Any], opcional) — Argumentos de palavra-chave adicionais passados para ~Trainer.create_model_card.

Envie self.model e self.processing_class para o hub de modelos do 🤗 no repositório self.args.hub_model_id.

GOLDConfig

classe trl.experimental.gold.GOLDConfig

< source >

( output_dir: str | None = Noneper_device_train_batch_size: int = 8num_train_epochs: float = 3.0max_steps: int = -1learning_rate: float = 1e-07lr_scheduler_type: transformers.trainer_utils.SchedulerType | str = 'linear'lr_scheduler_kwargs: dict | str | None = Nonewarmup_steps: float = 0optim: transformers.training_args.OptimizerNames | str = 'adamw_torch_fused'optim_args: str | None = Noneweight_decay: float = 0.0adam_beta1: float = 0.9adam_beta2: float = 0.999adam_epsilon: float = 1e-08optim_target_modules: None | str | list[str] = Nonegradient_accumulation_steps: int = 1average_tokens_across_devices: bool = Truemax_grad_norm: float = 1.0label_smoothing_factor: float = 0.0bf16: bool | None = Nonefp16: bool = Falsebf16_full_eval: bool = Falsefp16_full_eval: bool = Falsetf32: bool | None = Nonegradient_checkpointing: bool = Truegradient_checkpointing_kwargs: dict[str, typing.Any] | str | None = Nonetorch_compile: bool = Falsetorch_compile_backend: str | None = Nonetorch_compile_mode: str | None = Noneuse_liger_kernel: bool = Falseliger_kernel_config: dict[str, bool] | None = Noneuse_cache: bool = Falseneftune_noise_alpha: float | None = Nonetorch_empty_cache_steps: int | None = Noneauto_find_batch_size: bool = Falselogging_strategy: transformers.trainer_utils.IntervalStrategy | str = 'steps'logging_steps: float = 10logging_first_step: bool = Falselog_on_each_node: bool = Truelogging_nan_inf_filter: bool = Trueinclude_num_input_tokens_seen: str | bool = 'no'log_level: str = 'passive'log_level_replica: str = 'warning'disable_tqdm: bool | None = Nonereport_to: None | str | list[str] = 'none'run_name: str | None = Noneproject: str = 'huggingface'trackio_space_id: str | None = Nonetrackio_bucket_id: str | None = Nonetrackio_static_space_id: typing.Union[str, NoneType, typing.Literal[False]] = Noneeval_strategy: transformers.trainer_utils.IntervalStrategy | str = 'no'eval_steps: float | None = Noneeval_delay: float = 0per_device_eval_batch_size: int = 8prediction_loss_only: bool = Falseeval_on_start: bool = Falseeval_do_concat_batches: bool = Trueeval_use_gather_object: bool = Falseeval_accumulation_steps: int | None = Noneinclude_for_metrics: list[str] = <factory>batch_eval_metrics: bool = Falsesave_only_model: bool = Falsesave_strategy: transformers.trainer_utils.SaveStrategy | str = 'steps'save_steps: float = 500save_on_each_node: bool = Falsesave_total_limit: int | None = Noneenable_jit_checkpoint: bool = Falsepush_to_hub: bool = Falsehub_token: str | None = Nonehub_private_repo: bool | None = Nonehub_model_id: str | None = Nonehub_strategy: transformers.trainer_utils.HubStrategy | str = 'every_save'hub_always_push: bool = Falsehub_revision: str | None = Noneload_best_model_at_end: bool = Falsemetric_for_best_model: str | None = Nonegreater_is_better: bool | None = Noneignore_data_skip: bool = Falserestore_callback_states_from_checkpoint: bool = Falsefull_determinism: bool = Falseseed: int = 42data_seed: int | None = Noneuse_cpu: bool = Falseaccelerator_config: dict | str | None = Noneparallelism_config: accelerate.parallelism_config.ParallelismConfig | None = Nonedataloader_drop_last: bool = Falsedataloader_num_workers: int = 0dataloader_pin_memory: bool = Truedataloader_persistent_workers: bool = Falsedataloader_prefetch_factor: int | None = Nonedataloader_multiprocessing_context: str | None = Nonedataloader_in_order: bool = Trueremove_unused_columns: bool = Truelabel_names: list[str] | None = Nonetrain_sampling_strategy: str = 'random'length_column_name: str = 'length'ddp_find_unused_parameters: bool | None = Noneddp_bucket_cap_mb: int | None = Noneddp_broadcast_buffers: bool | None = Noneddp_static_graph: bool | None = Noneddp_backend: str | None = Noneddp_timeout: int = 1800fsdp: str | None = Nonefsdp_config: dict[str, typing.Any] | str | None = Nonedeepspeed: dict | str | None = Nonedebug: str | list[transformers.debug_utils.DebugOption] = ''skip_memory_metrics: bool = Truedo_train: bool = Falsedo_eval: bool = Falsedo_predict: bool = Falseresume_from_checkpoint: str | None = Nonelocal_rank: int = -1model_init_kwargs: dict[str, typing.Any] | str | None = Nonetrust_remote_code: bool = Falserouter_aux_loss_coef: float | None = Nonechat_template_path: str | None = Nonedataset_text_field: str = 'text'dataset_kwargs: dict[str, typing.Any] | str | None = Nonedataset_num_proc: int | None = Noneeos_token: str | None = Nonemax_length: int | None = 1024truncation_mode: str = 'keep_start'shuffle_dataset: bool = Falsepacking: bool = Falsepacking_strategy: str = 'bfd'padding_free: bool = Falsepad_to_multiple_of: int | str | None = 'auto'eval_packing: bool | None = Nonecompletion_only_loss: bool | None = Noneassistant_only_loss: bool = Falseloss_type: str = 'nll'activation_offloading: bool = Falsepad_token: str | None = Nonetemperature: float = 0.9top_p: float = 0.95top_k: int = 0lmbda: float = 0.5beta: float = 0.5max_completion_length: int = 128teacher_model_name_or_path: str | None = Noneteacher_model_revision: str | None = Noneteacher_model_init_kwargs: dict[str, typing.Any] | str | None = Noneteacher_tokenizer_name_or_path: str | None = Nonedisable_dropout: bool = Trueseq_kd: bool = Falsenum_generations: int = 1generation_batch_size: int | None = Noneuse_uld_loss: bool = Falseuld_token_merge_strategy: str = 'observed'use_extended_uld: bool = Trueuld_use_hybrid_loss: bool = Falseuld_hybrid_matched_weight: float | None = Noneuld_hybrid_unmatched_weight: float | None = Noneuld_crossentropy_weight: float = 0.0uld_distillation_weight: float = 1.0uld_student_temperature: float = 1.0uld_teacher_temperature: float = 1.0uld_skip_student_eos: bool = Trueuld_skip_teacher_eos: bool = Truextoken_loss_type: str = 'none'xtoken_projection_matrix_path: str | None = Nonextoken_temperature: float = 1.0xtoken_dynamic_scaling: bool = Truextoken_uncommon_topk: int = 8192xtoken_vocab_topk: int = 8192xtoken_kl_weight: float = 1.0xtoken_ce_scale: float = 0.1use_vllm: bool = Falsevllm_mode: str = 'colocate'vllm_server_base_url: str | None = Nonevllm_server_host: str = '0.0.0.0'vllm_server_port: int = 8001vllm_server_timeout: float = 240.0vllm_group_port: int = 51216vllm_gpu_memory_utilization: float = 0.9vllm_tensor_parallel_size: int = 1vllm_max_model_length: int | None = Nonevllm_model_impl: str = 'vllm'vllm_structured_outputs_regex: str | None = Nonevllm_sync_frequency: int = 1vllm_enable_sleep_mode: bool = Falselog_completions: bool = Falselog_completions_steps: int = 100num_completions_to_print: int | None = Nonewandb_log_unique_prompts: bool = Truecallbacks: list[str] = <factory> )

Parâmetros que controlam a geração e o loop de treinamento

temperature (float, opcional, padrão 0.9) — Temperatura para amostragem. Quanto maior a temperatura, mais aleatórias são as conclusões.
top_p (float, opcional, padrão 0.95) — Se definido para um float < 1, apenas o menor conjunto de tokens mais prováveis, cujas probabilidades somem top_p ou mais, é mantido para a geração.
top_k (int, opcional, padrão 0) — Número de tokens do vocabulário com maior probabilidade a serem mantidos para o filtro top-k. Se 0, o filtro top-k é desativado e todos os tokens são considerados.
lmbda (float, opcional, padrão 0.5) — Parâmetro Lambda que controla a fração de dados do aluno (ou seja, a proporção de saídas geradas pelo aluno on-policy).
beta (float, opcional, padrão 0.5) — Coeficiente de interpolação entre 0.0 e 1.0 da perda de Divergência Generalizada de Jensen-Shannon. Quando beta é 0.0, a perda é a divergência KL. Quando beta é 1.0, a perda é a Divergência KL Inversa.
max_completion_length (int, opcional, padrão 128) — Número máximo de tokens a serem gerados por conclusão.
teacher_model_name_or_path (str, opcional) — Nome ou caminho do modelo teacher. Se None, o modelo teacher será o mesmo que o modelo em treinamento.
teacher_model_revision (str ou None, opcional, padrão None) — Revisão do modelo do professor (por exemplo, nome da branch, tag ou hash de commit). Se None, a revisão padrão é utilizada.
teacher_model_init_kwargs (dict[str, Any], opcional) — Argumentos de palavra-chave para passar a AutoModelForCausalLM.from_pretrained ao instanciar o modelo teacher a partir de uma string. O valor de revision também é usado ao carregar a classe de processamento do teacher.
teacher_tokenizer_name_or_path (str, opcional) — Nome ou caminho do tokenizer para o modelo teacher. Se for None ao usar a perda ULD, será utilizado o mesmo tokenizer do modelo student (não recomendado para destilação entre tokenizers diferentes).
disable_dropout (bool, opcional, padrão True) — Se deve desabilitar o dropout no modelo.
seq_kd (bool, opcional, padrão False) — Parâmetro que controla se deve ser executado Sequence-Level KD (pode ser visto como FT supervisionado na saída gerada pelo teacher).
num_generations (int, opcional, padrão 1) — Número de gerações por prompt. Cada prompt é repetido esse número de vezes no lote de geração.
generation_batch_size (int ou None, opcional, padrão None) — Número de prompts únicos por worker por passo do otimizador. Se None, é calculado a partir de (per_device_train_batch_size * gradient_accumulation_steps) // num_generations.

Parâmetros que controlam a perda ULD

use_uld_loss (bool, opcional, padrão False) — Se deve usar a perda Universal Logit Distillation (ULD) em vez da perda Generalized Jensen-Shannon Divergence.
use_extended_uld (bool, opcional, padrão True) — Se deve habilitar o alinhamento estendido de ULD que usa tokenizers para alinhar e mesclar probabilidades de tokens entre as tokenizações do aluno e do professor. Quando True, o trainer calculará mapeamentos de tokens e mesclará probabilidades para tokens divididos; quando False, o ULD usará truncamento posicional simples, como no artigo original de ULD. Defina como False para tokenizers que não são nem ByteLevel nem BPE com fallback de bytes (por exemplo, Llama 2).
uld_token_merge_strategy (str, opcional, padrão "observed") — Estratégia usada para alinhar os logits de resposta e mesclar as probabilidades de tokens na perda ULD. Com "observed", os logits de resposta são fatiados nas posições da resposta e os tokens divididos (quando use_extended_uld=True) são mesclados multiplicando a distribuição marginal na primeira posição pelas probabilidades condicionais escalares dos tokens posteriores reais. Com "bayesian", o fatiamento dos logits de resposta é deslocado uma posição para trás, de modo que probs[k] preveja token_ids[k] (regra da cadeia), e os tokens divididos são mesclados usando a distribuição completa da última posição, condicionada aos tokens de prefixo reais, multiplicada pelas probabilidades escalares dos tokens anteriores. O deslocamento do alinhamento de logits se aplica independentemente de use_extended_uld estar habilitado.
uld_use_hybrid_loss (bool, opcional, padrão False) — Se deve usar uma perda híbrida que combina a perda ULD e a perda JSD. Quando True, a perda final é uma combinação de JSD para mapeamentos de tokens conhecidos e ULD para mapeamentos de tokens desconhecidos.
uld_hybrid_matched_weight (float or None, opcional) — Peso para o componente de perda de tokens correspondentes ao usar a perda híbrida ULD + JSD. Este peso escala a perda JSD calculada sobre tokens que possuem um mapeamento direto entre as tokenizações do aluno e do professor. Se None, usa ponderação adaptativa baseada na sobreposição de vocabulário. Deve ser definido junto com uld_hybrid_unmatched_weight (ambos None ou ambos float).
uld_hybrid_unmatched_weight (float ou None, opcional) — Peso para o componente de perda de tokens não correspondidos ao usar a perda híbrida ULD + JSD. Este peso escala a perda ULD calculada sobre tokens que não possuem um mapeamento direto entre as tokenizações do aluno e do professor. Se None, usa ponderação adaptativa com base na sobreposição de vocabulário. Deve ser definido junto com uld_hybrid_matched_weight (ambos None ou ambos float).
uld_crossentropy_weight (float, opcional, padrão 0.0) — Peso para o componente de perda de entropia cruzada na perda ULD. Se for 0, apenas a perda de destilação ULD é usada.
uld_distillation_weight (float, opcional, padrão 1.0) — Peso para o componente de perda de destilação na perda ULD.
uld_student_temperature (float, opcional, padrão 1.0) — Temperatura para os logits do aluno no cálculo da perda ULD.
uld_teacher_temperature (float, opcional, padrão 1.0) — Temperatura para os logits do professor no cálculo da perda ULD.
uld_skip_student_eos (bool, opcional, padrão True) — Se deve ignorar o token EOS para o aluno no cálculo da perda ULD.
uld_skip_teacher_eos (bool, opcional, padrão True) — Se deve ignorar o token EOS do teacher no cálculo da perda ULD.

Parâmetros para KD de cross-tokenizer X-Token (https

xtoken_loss_type (str, opcional, padrão "none") — Variante de KD entre tokenizadores: "none" (desativado), "p_kl" (KL de projeção) ou "h_kl" (KL híbrido com conjunto comum relaxado). Exige teacher_tokenizer_name_or_path e xtoken_projection_matrix_path.
xtoken_projection_matrix_path (str, opcional) — Caminho para um arquivo de matriz de projeção pré-computada. Obrigatório quando xtoken_loss_type != "none".
xtoken_temperature (float, opcional, padrão 1.0) — Temperatura T para a perda KD de X-Token; a perda é multiplicada por T² (Hinton 2015).
xtoken_dynamic_scaling (bool, opcional, padrão True) — Escala o KD por stop_gradient(abs(CE) / abs(KD)) em cada microbatch para equilibrar as magnitudes das perdas (Eq. 7 do paper).
xtoken_uncommon_topk (int, opcional, padrão 8192) — H-KL: limita a comparação de L1 não comum ordenada aos top-k tokens de cada lado.
xtoken_vocab_topk (int, opcional, padrão 8192) — P-KL: restringe o KL aos top-k tokens globais do teacher com o maior logit. Interage com a matriz de projeção --top-k: tokens do teacher fora do alcance da matriz recebem massa projetada do student próxima de zero e dominam o KL, portanto, valores menores concentram a perda nos tokens que o student consegue corresponder.
xtoken_kl_weight (float, opcional, padrão 1.0) — Peso para o termo KD de X-Token quando xtoken_dynamic_scaling=False.
xtoken_ce_scale (float, opcional, padrão 0.1) — Fator de escala para o termo de CE na perda X-Token quando xtoken_dynamic_scaling=False.

Parâmetros que controlam a integração com o vLLM

use_vllm (bool, opcional, padrão False) — Se deve usar o vLLM para gerar conclusões a partir do modelo aluno. Exige que o vllm esteja instalado.
vllm_mode (str, opcional, padrão "colocate") — Modo para a integração do vLLM do aluno. Pode ser "server" (conectar a um servidor vLLM em execução) ou "colocate" (executar o vLLM no mesmo processo).
vllm_server_host (str, opcional, padrão "0.0.0.0") — Host do servidor vLLM para o modelo aluno (se vllm_mode="server").
vllm_server_port (int, opcional, padrão 8001) — Porta do servidor vLLM para o modelo aluno (se vllm_mode="server").
vllm_server_timeout (float, opcional, padrão 240.0) — Tempo limite para conectar ao servidor vLLM do aluno (se vllm_mode="server").
vllm_gpu_memory_utilization (float, opcional, padrão 0.9) — Utilização da memória de GPU para o motor vLLM do aluno em modo colocalizado (se vllm_mode="colocate"). É recomendado definir este valor para um número baixo se os modelos de aluno e professor compartilharem a mesma GPU.
vllm_tensor_parallel_size (int, opcional, padrão 1) — Tamanho do paralelismo de tensor para o motor vLLM do aluno co-localizado (se vllm_mode="colocate").
vllm_structured_outputs_regex (str, opcional) — Regex para saídas estruturadas do vLLM para o modelo aluno.
vllm_server_base_url (str, opcional) — URL base para o servidor vLLM (por exemplo, "http://localhost:8001"). Se fornecida, vllm_server_host e vllm_server_port são ignorados.
vllm_group_port (int, opcional, padrão 51216) — Porta para o grupo de atualização de pesos do vLLM (comunicador NCCL). A menos que a porta esteja ocupada, não é necessário alterá-la.
vllm_max_model_length (int, opcional) — Comprimento máximo da sequência do modelo para o mecanismo vLLM co-localizado quando vllm_mode="colocate". O valor padrão é o comprimento máximo de contexto do modelo.
vllm_model_impl (str, opcional, padrão "vllm") — Backend de implementação do modelo a ser usado no vLLM. Use "vllm" (padrão) ou "transformers".
vllm_sync_frequency (int, opcional, padrão 1) — Frequência (em passos de treinamento) para sincronizar os pesos do modelo aluno com o motor vLLM. Defina como 1 para sincronizar após cada passo.
vllm_enable_sleep_mode (bool, opcional, padrão False) — Ativa o modo de suspensão do vLLM para descarregar os pesos/caches do aluno durante o passo do otimizador. Mantém o uso de memória da GPU baixo, mas a reativação do motor adiciona latência de transferência entre host e dispositivo.

Parâmetros que controlam o registro de logs

log_completions (bool, opcional, padrão False) — Se deve registrar uma amostra de pares (prompt, completion) a cada logging_steps passos. Se o rich estiver instalado, ele imprime a amostra. Se o registro do wandb estiver habilitado, ele registra a amostra no wandb.
log_completions_steps (int, opcional, padrão 100) — Número de passos entre o registro de pares (prompt, completion). Usado apenas se log_completions estiver definido como True.
num_completions_to_print (int ou None, opcional) — Número de conclusões a serem exibidas com rich. Se None, todas as conclusões são registradas.
wandb_log_unique_prompts (bool, opcional, padrão True) — Se os prompts únicos devem ser registrados no wandb. Isso criará uma nova execução para cada prompt único.
callbacks (list[str], opcional, padrão []) — Os callbacks a serem executados durante o treinamento.

Classe de configuração para experimental.gold.GOLDTrainer.

Esta classe inclui apenas os parâmetros específicos para o treinamento GOLD. Para a lista completa de argumentos de treinamento, consulte a documentação de TrainingArguments e SFTConfig.

Esses parâmetros possuem valores padrão diferentes de TrainingArguments:

learning_rate: O valor padrão é 1e-7 em vez de 5e-5.

Update on GitHub