TRL documentation

Distillation Trainer

TRL

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v1.17.1).

Join the Hugging Face community

and get access to the augmented documentation experience

Collaborate on models, datasets and Spaces

Faster examples with accelerated inference

Switch between documentation themes

to get started

Distillation Trainer

Visão geral

O Distillation Trainer implementa a destilação de conhecimento on-policy conforme descrito em On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes, de Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos, Matthieu Geist e Olivier Bachem.

O resumo do artigo é o seguinte:

A destilação de conhecimento (KD) é amplamente utilizada para comprimir um modelo professor, reduzindo seu custo de inferência e sua pegada de memória, ao treinar um modelo aluno menor. No entanto, os métodos atuais de KD para modelos de sequência auto-regressivos sofrem com o desalinhamento de distribuição entre as sequências de saída vistas durante o treinamento e aquelas geradas pelo aluno durante a inferência. Para abordar essa questão, introduzimos a Destilação de Conhecimento Generalizada (GKD). Em vez de depender exclusivamente de um conjunto fixo de sequências de saída, o GKD treina o aluno em suas próprias sequências de saída auto-geradas, aproveitando o feedback do professor sobre tais sequências. Diferente das abordagens de KD supervisionadas, o GKD também oferece a flexibilidade de empregar funções de perda alternativas entre o aluno e o professor, o que pode ser útil quando o aluno carece da expressividade para imitar a distribuição do professor.

O DistillationTrainer treina um modelo aluno menor para corresponder à distribuição de próximo token de um professor nas próprias gerações on-policy do aluno. Ele gera as conclusões do aluno on-policy (opcionalmente com vLLM) e corresponde o professor sobre sua distribuição completa de próximo token com uma perda de divergência de Jensen-Shannon em blocos eficiente em memória, de modo que a distribuição densa do professor nunca é materializada por completo.

Este trainer foi contribuído por Carlos Miguel Patiño.

Para uma contraparte assíncrona que desacopla a geração do treinamento e calcula as pontuações em relação a um professor servido via HTTP, em vez de uma passagem direta local, consulte AsyncDistillationTrainer.

Início rápido

Este exemplo demonstra como treinar um modelo usando o método de destilação. Destilamos um modelo Qwen 2.5 0.5B Instruct a partir de um modelo professor Qwen 2.5 1.5B Instruct nos prompts do conjunto de dados de prompts UltraFeedback. Você pode visualizar os dados do conjunto de dados aqui:

Abaixo está o script para treinar o modelo.

# train_distillation.py
from datasets import load_dataset
from trl import DistillationTrainer

dataset = load_dataset("trl-lib/ultrafeedback-prompt", split="train")

trainer = DistillationTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",
    teacher_model="Qwen/Qwen2.5-1.5B-Instruct",
    train_dataset=dataset,
)
trainer.train()

Execute o script usando o seguinte comando:

accelerate launch train_distillation.py

Olhando mais a fundo para o método de destilação

A destilação de conhecimento on-policy treina um aluno para reproduzir a distribuição de próximo token de um professor em conclusões geradas pelo próprio aluno, em vez de em um conjunto fixo de saídas do professor. Aprender a partir de suas próprias gerações permite que o aluno corrija seus próprios erros, o que geralmente supera a destilação off-policy. Esta seção detalha como isso funciona na prática, cobrindo as duas etapas principais: gerar conclusões e calcular a perda.

Gerando complementos

Em cada etapa de treinamento, o aluno gera um lote de conclusões para os prompts amostrados.

Calculando a perda

A perda é a divergência de Jensen-Shannon generalizada (JSD) entre a distribuição do aluno $p_S$ e a distribuição do professor $p_T$ sobre os tokens de conclusão gerados, interpolada por beta e definida como: $\mathcal{L}_\beta = \beta \, \mathbb{D}_{\mathrm{KL}}\!\left[ p_T \| p_M \right] + (1 - \beta) \, \mathbb{D}_{\mathrm{KL}}\!\left[ p_S \| p_M \right], \qquad p_M = (1 - \beta) \, p_S + \beta \, p_T,$

onde $p_M$ é a $\beta$ -mistura das duas distribuições. Os extremos reduzem-se às divergências puras: beta=0.0 produz o KL direto $\mathbb{D}_{\mathrm{KL}}\!\left[ p_T \| p_S \right]$ e beta=1.0 o KL reverso $\mathbb{D}_{\mathrm{KL}}\!\left[ p_S \| p_T \right]$ .

Na prática, a projeção para os logits de vocabulário e a divergência são calculados em blocos, de modo que a memória de ativação de pico não escala com o tensor completo de logits de vocabulário × comprimento da sequência. Consulte Reduzindo o uso de memória.

Tipo de dataset esperado

O conjunto de dados deve estar formatado como um conjunto de dados conversacional somente com prompt. O aluno gera seus próprios completamentos on-policy, portanto, apenas o prompt é necessário:

{"prompt": [{"role": "user", "content": "What color is the sky?"}]}

Métricas registradas

Durante o treinamento e a avaliação, registramos as seguintes métricas:

num_tokens: O número total de tokens processados até o momento, incluindo tanto prompts quanto completions. Ao usar ferramentas, apenas tokens que não são de ferramentas são contabilizados.
step_time: O tempo médio (em segundos) por etapa de treinamento (incluindo a geração).
completions/mean_length: O comprimento médio das conclusões geradas. Ao usar ferramentas, apenas os tokens que não são de ferramentas são contados.
completions/min_length: O comprimento mínimo das conclusões geradas. Ao usar ferramentas, apenas os tokens que não são de ferramentas são contados.
completions/max_length: O comprimento máximo das conclusões geradas. Ao usar ferramentas, apenas os tokens que não são de ferramentas são contados.
completions/mean_terminated_length: O comprimento médio das conclusões geradas que terminam com EOS. Ao usar ferramentas, apenas os tokens que não são de ferramentas são contados.
completions/min_terminated_length: O comprimento mínimo das conclusões geradas que terminam com EOS. Ao usar ferramentas, apenas os tokens que não são de ferramentas são contados.
completions/max_terminated_length: O comprimento máximo das completions geradas que terminam com EOS. Ao usar ferramentas, apenas os tokens que não são de ferramentas são contados.
completions/clipped_ratio: A proporção de conclusões truncadas (clipped).
tools/call_frequency: O número médio de chamadas de ferramentas por conclusão no lote de geração. Registrado apenas quando tools são fornecidas.
tools/failure_frequency: A fração de chamadas de ferramenta que falharam (a ferramenta não foi encontrada, lançou uma exceção ou o tipo de chamada não é suportado). É 0.0 quando nenhuma ferramenta foi chamada. Registrado apenas quando tools são fornecidas.
entropy: Entropia média das previsões de tokens entre as conclusões geradas (em nats).

Personalização

Acelere o treinamento com geração impulsionada por vLLM

A geração é frequentemente o principal gargalo ao treinar com métodos on-policy. Para acelerar a geração, você pode usar vLLM, um motor de inferência de alto throughput e baixa latência para LLMs. Para ativá-lo, primeiro instale o pacote com

pip install trl[vllm]

Suportamos duas formas de usar o vLLM durante o treinamento: modo colocalizado e modo servidor.

Opção 1: Modo Colocate

Neste modo, o vLLM é executado dentro do processo do trainer e compartilha a memória da GPU com o modelo de treinamento. Isso evita a inicialização de um servidor separado e pode melhorar a utilização da GPU, mas pode levar a disputas de memória nas GPUs de treinamento. Este é o modo padrão.

from trl import DistillationConfig

training_args = DistillationConfig(
    ...,
    use_vllm=True,  # vllm_mode="colocate" by default
)

Opção 2: Modo servidor

Neste modo, o vLLM é executado em um processo separado (e usando GPUs separadas) e se comunica com o treinador via HTTP. Isso é ideal se você tiver GPUs dedicadas para inferência.

Inicie o servidor vLLM:

VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \
    --weight-transfer-config '{"backend": "nccl"}' \
    --logprobs-mode processed_logprobs \
    --max-logprobs -1

Ative o modo de servidor no seu script de treinamento:

from trl import DistillationConfig

training_args = DistillationConfig(
    ...,
    use_vllm=True,
    vllm_mode="server",
)

Certifique-se de que o servidor está usando GPUs diferentes das do trainer; caso contrário, você pode encontrar erros do NCCL. Você pode especificar as GPUs a serem usadas com a variável de ambiente CUDA_VISIBLE_DEVICES.

Dependendo do tamanho do modelo e dos requisitos gerais de memória de GPU para o treinamento, você pode precisar ajustar o parâmetro vllm_gpu_memory_utilization em DistillationConfig para evitar subutilização ou erros de memória insuficiente.

Para mais informações, consulte Aceleração do treinamento com vLLM.

Treine adaptadores com PEFT

Suportamos integração estreita com a biblioteca 🤗 PEFT, permitindo que você treine adaptadores e os compartilhe no Hub, em vez de treinar o aluno inteiro.

from datasets import load_dataset
from trl import DistillationTrainer
from peft import LoraConfig

dataset = load_dataset("trl-lib/ultrafeedback-prompt", split="train")

trainer = DistillationTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",
    teacher_model="Qwen/Qwen2.5-1.5B-Instruct",
    train_dataset=dataset,
    peft_config=LoraConfig(),
)
trainer.train()

A perda de destilação lê lm_head.weight diretamente e executa o backbone do aluno sem passar por PeftModel.forward(). Adapters em lm_head (via target_modules) e métodos de prompt-learning (PromptTuning, PrefixTuning, P-Tuning) são, portanto, rejeitados, pois seriam ignorados silenciosamente. Para treinar o head, use modules_to_save=["lm_head"] em vez disso.

Treinamento de Agentes

DistillationTrainer suporta treinamento de agentes: o aluno chama ferramentas durante a geração e é destilado na trajetória resultante. Os tokens de resultado das ferramentas são mascarados na perda, de modo que o aluno é treinado apenas nos tokens que ele mesmo gerou.

Ferramentas

O argumento tools espera uma lista de funções Python que definem as ferramentas disponíveis para o agente:

from trl import DistillationTrainer

trainer = DistillationTrainer(
    tools=[tool1, tool2],
    ...,
)

Cada ferramenta deve ser uma função Python padrão com argumentos e tipos de retorno com type hints, além de uma docstring no estilo Google descrevendo seu propósito, argumentos e valor de retorno. Para mais detalhes, consulte o Guia de passagem de ferramentas.

O loop de chamada de ferramentas exige que o modelo de chat seja preservador de prefixo (a adição de uma mensagem de ferramenta não deve alterar a forma como as mensagens anteriores são renderizadas). Para famílias de modelos conhecidas (por exemplo, Qwen3, DeepSeek-V3), o TRL substitui automaticamente o modelo de treinamento por um modelo corrigido quando as ferramentas estão habilitadas. Consulte Modelos de Chat para a lista completa.

Use max_tool_calling_iterations no DistillationConfig para limitar o número de turnos de chamada de ferramentas. Por padrão, não há limite, e a geração é interrompida quando o aluno produz um turno de resposta sem chamadas de ferramentas.

Exemplo:

from trl import DistillationTrainer

def multiply(a: int, b: int) -> int:
    """
    Multiplies two integers.

    Args:
        a: The first integer.
        b: The second integer.

    Returns:
        The product of the two integers.
    """
    return a * b

trainer = DistillationTrainer(
    tools=[multiply],
    ...,
)

Ferramentas assíncronas ainda não são suportadas — passe funções síncronas.

Respostas de Ferramentas Multimodais

Ferramentas podem retornar imagens junto com texto ao retornar uma lista de blocos de conteúdo. Isso é útil para o treinamento de agentes VLM em que a ferramenta fornece feedback visual (por exemplo, capturas de tela, gráficos, capturas de câmera).

from PIL import Image

def take_screenshot() -> list:
    """
    Takes a screenshot of the current screen.

    Returns:
        The screenshot image with a description.
    """
    img = Image.open("screenshot.png")
    return [{"type": "image", "image": img}, {"type": "text", "text": "Here is the screenshot."}]

As imagens retornadas são injetadas automaticamente na conversa e passadas ao VLM para as subsequentes rodadas de geração.

Treinando Modelos de Visão e Linguagem

O DistillationTrainer suporta a destilação de Vision-Language Models (VLMs) em datasets multimodais que contêm tanto texto quanto imagens. Passe um VLM como aluno e como professor, e forneça um dataset apenas com prompts, contendo uma coluna image (uma imagem por amostra) ou uma coluna images (lista de imagens por amostra). Para mais informações sobre a estrutura esperada do dataset, consulte a seção Dataset Format — Vision datasets.

Testado com:

Gemma 3 — por exemplo, google/gemma-3-4b-it
LLaVA-NeXT — por exemplo, llava-hf/llava-v1.6-mistral-7b-hf
Qwen2-VL — por exemplo, Qwen/Qwen2-VL-2B-Instruct
Qwen2.5-VL — por exemplo, Qwen/Qwen2.5-VL-3B-Instruct

A compatibilidade com todos os VLMs não é garantida. Se você acredita que um modelo deve ser suportado, sinta-se à vontade para abrir uma issue no GitHub — ou, melhor ainda, enviar um pull request com as alterações necessárias.

Interface de linha de comando

Use o trl distillation CLI para iniciar o treinamento de destilação pela linha de comando. Ele suporta treinamento completo e LoRA por meio das flags padrão de ModelConfig.

# Full training:
trl distillation \
    --model_name_or_path Qwen/Qwen2.5-0.5B-Instruct \
    --teacher_model_name_or_path Qwen/Qwen2.5-1.5B-Instruct \
    --dataset_name trl-lib/ultrafeedback-prompt \
    --learning_rate 2e-5 \
    --per_device_train_batch_size 4 \
    --gradient_accumulation_steps 8 \
    --output_dir distilled-model \
    --num_train_epochs 1

# LoRA:
trl distillation \
    --model_name_or_path Qwen/Qwen2.5-0.5B-Instruct \
    --teacher_model_name_or_path Qwen/Qwen2.5-1.5B-Instruct \
    --dataset_name trl-lib/ultrafeedback-prompt \
    --learning_rate 2e-4 \
    --per_device_train_batch_size 4 \
    --gradient_accumulation_steps 8 \
    --output_dir distilled-model \
    --num_train_epochs 1 \
    --use_peft \
    --lora_r 64 \
    --lora_alpha 16

DistillationTrainer

classe trl.DistillationTrainer

< source >

( model: str | PreTrainedModel | PeftModelteacher_model: str | transformers.modeling_utils.PreTrainedModel | None = Noneargs: trl.trainer.distillation_config.DistillationConfig | 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.processing_utils.ProcessorMixin | None = Nonecallbacks: list[transformers.trainer_callback.TrainerCallback] | None = Noneoptimizers: tuple[typing.Optional[torch.optim.Optimizer], typing.Optional[torch.optim.lr_scheduler.LambdaLR]] = (None, None)quantization_config: BitsAndBytesConfig | None = Nonepeft_config: PeftConfig | None = Nonetools: list[collections.abc.Callable] | None = None )

Parâmetros

model (str ou PreTrainedModel ou PeftModel) — Modelo a ser treinado. Pode ser:
- Uma string, sendo o model id de um modelo pré-treinado hospedado em um repositório de modelos no huggingface.co, ou um caminho para um diretório contendo pesos de modelo salvos usando save_pretrained, por exemplo, './my_model_directory/'. O modelo é carregado usando <ModelArchitecture>.from_pretrained (onde <ModelArchitecture> é derivado da configuração do modelo) com os argumentos de palavra-chave em args.model_init_kwargs. Se dtype não for especificado em args.model_init_kwargs, o padrão é float32. Isso difere de from_pretrained, onde (desde a versão 5 do Transformers) o dtype é inferido a partir da configuração do modelo.
- Um objeto PreTrainedModel. Somente modelos de linguagem causais são suportados.
- Um objeto PeftModel. Somente modelos de linguagem causais são suportados.
teacher_model (str ou PreTrainedModel, opcional) — Modelo de professor cuja distribuição de próximo token o aluno é treinado para corresponder. Pode ser um id de modelo / caminho (carregado como model, usando args.teacher_model_init_kwargs) ou um PreTrainedModel instanciado. Ele deve compartilhar o vocabulário do aluno. Pode ser omitido por subclasses que fornecem o professor de outra forma (por exemplo, um servidor remoto).
args (DistillationConfig, opcional) — Configuração para este trainer. Se None, uma configuração padrão é usada.
train_dataset (Dataset ou IterableDataset) — Dataset a ser usado para o treinamento. Ele deve incluir uma coluna "prompt". Colunas adicionais no dataset são ignoradas. O formato das amostras pode ser:
- Padrão: Cada amostra contém texto simples.
- Conversacional: cada amostra contém mensagens estruturadas (por exemplo, papel e conteúdo).
Quando train_dataset é um IterableDataset (por exemplo, um dataset em streaming), max_steps deve ser definido nos argumentos de treinamento, pois seu comprimento não pode ser inferido e o número total de etapas de treinamento é necessário para limitar o loop de treinamento e configurar o agendador de taxa de aprendizado.
eval_dataset (Dataset, IterableDataset, DatasetDict, IterableDatasetDict ou dict[str, Dataset | IterableDataset]) — Dataset a ser usado para avaliação. Ele deve atender aos mesmos requisitos que train_dataset.
processing_class (PreTrainedTokenizerBase, ProcessorMixin, opcional) — Classe de processamento usada para processar os dados. O lado do padding deve ser definido como “left”. Se None, a classe de processamento é carregada a partir do nome do modelo com from_pretrained. Um token de padding, tokenizer.pad_token, deve ser definido. Se a classe de processamento não tiver definido um token de padding, tokenizer.eos_token será usado como padrão.
callbacks (lista de TrainerCallback, opcional) — Lista de callbacks para personalizar o loop de treinamento. Adicionará esses à lista de callbacks padrão detalhada aqui.
Se você deseja remover um dos callbacks padrão utilizados, use o método remove_callback.
optimizers (tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None], opcional, padrão (None, None)) — Uma tupla contendo o otimizador e o agendador a serem usados. Por padrão, será uma instância de AdamW no seu modelo e um agendador dado por get_linear_schedule_with_warmup controlado por args.
quantization_config (BitsAndBytesConfig, opcional) — Configuração de quantização usada ao carregar o modelo a partir de um identificador de modelo. Combine com peft_config para treinamento QLoRA. Ignorado se o modelo já estiver instanciado.
peft_config (PeftConfig, opcional) — Configuração PEFT usada para envolver o modelo. Se None, o modelo não é envolvido.
tools (lista de Callable, opcional) — Uma lista de funções de ferramenta chamáveis que o modelo pode invocar durante a geração. Cada ferramenta deve ser uma função Python padrão com argumentos e valores de retorno com type hints apropriados, e um docstring no estilo Google descrevendo seu propósito, argumentos e valor de retorno. Para mais detalhes, veja: https://huggingface.co/docs/transformers/en/chat_extras#passing-tools. O modelo usa o nome da função, os type hints e o docstring para determinar como chamá-la. Certifique-se de que o chat template do modelo suporta o uso de ferramentas e de que ele foi fine-tuned para tool calling.

Trainer para destilação de conhecimento. O aluno é treinado on-policy — ele mesmo gera as conclusões — para corresponder à distribuição de próximo token do professor sob uma divergência de Jensen-Shannon generalizada (interpolando KL direta, KL reversa e JSD via beta), conforme introduzido em On-Policy Distillation of Language Models.

Exemplo:

>>> from trl import DistillationTrainer
>>> from datasets import load_dataset

>>> dataset = load_dataset("trl-lib/tldr", split="train")

>>> trainer = DistillationTrainer(
...     model="Qwen/Qwen2.5-0.5B-Instruct",
...     teacher_model="Qwen/Qwen2.5-1.5B-Instruct",
...     train_dataset=dataset,
... )
>>> trainer.train()

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.

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.

DistillationConfig

classe trl.DistillationConfig

< source >

( output_dir: str | None = Noneper_device_train_batch_size: int = 8num_train_epochs: float = 3.0max_steps: int = -1learning_rate: float = 1e-06lr_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 | None = Falselabel_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 = Falseteacher_model_name_or_path: str | None = Noneteacher_model_revision: str | None = Noneteacher_model_init_kwargs: dict[str, typing.Any] | str | None = Nonedisable_dropout: bool = Falsemax_completion_length: int | None = 512ds3_gather_for_generation: bool = Trueshuffle_dataset: bool | None = Truepad_to_multiple_of: int | str | None = 'auto'temperature: float = 1.0top_p: float = 1.0top_k: int = 0min_p: float | None = Nonegeneration_kwargs: dict | str | None = Nonechat_template_kwargs: dict | str | None = Nonerepetition_penalty: float = 1.0cache_implementation: str | None = Noneuse_vllm: bool = Falsevllm_mode: str = 'colocate'vllm_model_impl: str = 'vllm'vllm_enable_sleep_mode: bool = Falsevllm_structured_outputs_regex: str | None = Nonevllm_server_base_url: str | None = Nonevllm_server_host: str = '0.0.0.0'vllm_server_port: int = 8000vllm_server_timeout: float = 240.0vllm_group_port: int = 51216vllm_gpu_memory_utilization: float = 0.3vllm_max_model_length: int | None = Nonevllm_tensor_parallel_size: int = 1beta: float = 1.0max_tool_calling_iterations: int | None = Nonelog_completions: bool = Falsenum_completions_to_print: int | None = Nonelog_unique_prompts: bool = False )

Parâmetros que controlam o modelo e o modelo teacher

model_init_kwargs (str ou dict[str, Any], opcional) — Argumentos de palavra-chave para AutoModelForCausalLM.from_pretrained, usados quando o argumento model do trainer é fornecido como uma string. O valor de revision também é usado ao carregar a classe de processamento.
trust_remote_code (bool, opcional, padrão False) — Se permite carregar modelos e tokenizers que incluem código Python personalizado do Hub. Encaminhado para from_pretrained e from_pretrained, para o aluno e o professor.
teacher_model_name_or_path (str, opcional) — Nome ou caminho do modelo para o modelo teacher. Usado quando o teacher é carregado localmente.
teacher_model_revision (str, opcional) — Revisão do modelo do professor (por exemplo, nome da branch, tag ou hash do commit).
teacher_model_init_kwargs (str ou dict[str, Any], opcional) — Argumentos de palavra-chave passados para AutoModelForCausalLM.from_pretrained ao instanciar o modelo teacher a partir de uma string.
disable_dropout (bool, opcional, padrão False) — Se deve desabilitar o dropout no modelo aluno durante o treinamento.

Parâmetros que controlam o pré-processamento dos dados

remove_unused_columns (bool, opcional, padrão False) — Se deve manter apenas a coluna "prompt" no dataset. O trainer consome a coluna de prompt bruto e gera completions on-policy, portanto o padrão é False.
max_completion_length (int ou None, opcional, padrão 512) — Número máximo de tokens a serem gerados por conclusão durante a geração on-policy.
ds3_gather_for_generation (bool, opcional, padrão True) — Esta configuração se aplica ao DeepSpeed ZeRO-3. Se habilitado, os pesos do modelo de política são reunidos para a geração, melhorando a velocidade de geração. No entanto, desabilitar esta opção permite treinar modelos que excedem a capacidade de VRAM de uma única GPU, embora com o custo de uma geração mais lenta. Desabilitar esta opção não é compatível com a geração do vLLM.
shuffle_dataset (bool, opcional, padrão True) — Se deve embaralhar o conjunto de dados de treinamento.
pad_to_multiple_of (int ou str, opcional, padrão "auto") — Se definido, os ids dos prompts e dos completions serão preenchidos até um múltiplo deste valor. "auto" preenche até um múltiplo de 128 no Neuron (AWS Trainium), que compila uma vez por forma de entrada, e não preenche em outros ambientes.

Parâmetros que controlam a geração

temperature (float, opcional, padrão 1.0) — Temperatura para amostragem durante a geração e para o cálculo da perda de destilação. Valores maiores produzem distribuições de probabilidade mais suaves.
top_p (float, opcional, padrão 1.0) — Parâmetro de amostragem Top-p (núcleo) para geração on-policy.
top_k (int, opcional, padrão 0) — Parâmetro de amostragem top-k para geração on-policy. 0 desativa o filtro top-k.
min_p (float, opcional) — Probabilidade mínima do token, que será escalada pela probabilidade do token mais provável. Deve ser um valor entre 0.0 e 1.0. Valores típicos estão na faixa de 0.01-0.2.
generation_kwargs (dict[str, Any], opcional) — Argumentos de palavra-chave adicionais para passar ao GenerationConfig (se estiver usando transformers) ou SamplingParams (se estiver usando vLLM) ao amostrar as conclusões. Isso pode ser usado para personalizar ainda mais o comportamento da geração, como definir suppress_tokens, num_beams, etc. Se contiver chaves que conflitam com os outros parâmetros de geração (como min_p, top_p, etc.), eles os substituirão.
chat_template_kwargs (dict[str, Any], opcional) — Argumentos de palavra-chave adicionais para passar à função apply_chat_template ao gerar completions.
repetition_penalty (float, opcional, padrão 1.0) — Float que penaliza novos tokens com base em sua presença no prompt e no texto gerado até o momento. Valores > 1.0 incentivam o modelo a usar novos tokens, enquanto valores < 1.0 incentivam o modelo a repetir tokens.
cache_implementation (str, opcional) — Implementação do método de cache para uma geração mais rápida quando use_vllm está definido como False.

Parâmetros que controlam a aceleração da geração com vLLM

use_vllm (bool, opcional, padrão False) — Se deve usar o vLLM para gerar completions on-policy a partir do modelo aluno.
vllm_mode (str, opcional, padrão "colocate") — Modo para a integração do vLLM do aluno. Pode ser "server" ou "colocate".
vllm_model_impl (str, opcional, padrão "vllm") — Backend de implementação do modelo para vLLM. Use "vllm" ou "transformers".
vllm_enable_sleep_mode (bool, opcional, padrão False) — Ativa o modo de suspensão do vLLM para descarregar os pesos do aluno durante o passo do otimizador.
vllm_structured_outputs_regex (str, opcional) — Padrão de expressão regular para saídas estruturadas do vLLM.

Parâmetros que controlam o servidor vLLM (usados apenas quando `vllm_mode` é `"server"`)

vllm_server_base_url (str, opcional) — URL base para o servidor vLLM do aluno. Se fornecida, vllm_server_host e vllm_server_port são ignorados.
vllm_server_host (str, opcional, padrão "0.0.0.0") — Host do servidor vLLM do aluno.
vllm_server_port (int, opcional, padrão 8000) — Porta do servidor vLLM do aluno.
vllm_server_timeout (float, opcional, padrão 240.0) — Tempo limite para conectar ao servidor vLLM do aluno.
vllm_group_port (int, opcional, padrão 51216) — Porta para o grupo de atualização de pesos do vLLM (comunicador NCCL).

Parâmetros que controlam a execução do vLLM co-localizado (usados apenas quando `vllm_mode` é `"colocate"`)

vllm_gpu_memory_utilization (float, opcional, padrão 0.3) — Utilização da memória de GPU para o motor vLLM do aluno co-localizado.
vllm_max_model_length (int, opcional) — Comprimento máximo da sequência do modelo para o mecanismo vLLM co-localizado.
vllm_tensor_parallel_size (int, opcional, padrão 1) — Tamanho do paralelismo de tensor para o motor vLLM do aluno co-localizado.

Parâmetros que controlam o treinamento

beta (float, opcional, padrão 1.0) — Coeficiente de interpolação para a perda de Divergência Generalizada de Jensen-Shannon. Quando 0.0, a perda é a divergência KL direta. Quando 1.0, a perda é a divergência KL reversa. Quando 0.5, é a JSD padrão. Diferente do beta do GRPO (um coeficiente de penalidade KL em relação a um modelo de referência), aqui ele seleciona a própria divergência; não há penalidade KL de modelo de referência.
max_tool_calling_iterations (int, opcional) — Número máximo de turnos de chamada de ferramentas ao treinar um agente. Se None, não há limite e a geração para quando o modelo gera um turno de resposta sem chamadas de ferramentas ou quando o comprimento total da resposta atinge max_model_length.

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 e/ou do trackio estiver habilitado, ele registra a amostra no wandb e/ou no trackio.
num_completions_to_print (int, opcional) — Número de conclusões a serem exibidas com rich. Se None, todas as conclusões são registradas.
log_unique_prompts (bool, opcional, padrão False) — Se deve registrar prompts únicos. Se True, apenas prompts únicos são registrados. Se False, todos os prompts são registrados.

Classe de configuração para o DistillationTrainer.

Estende TrainingArguments com parâmetros específicos para destilação de conhecimento. Todos os campos necessários são declarados aqui.

Usando HfArgumentParser, podemos transformar essa classe em argumentos argparse que podem ser especificados na linha de comando.

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

logging_steps: O valor padrão é 10 em vez de 500.

gradient_checkpointing: O padrão é True em vez de False.

bf16: Padrão é True se fp16 não estiver definido, em vez de False.

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

Update on GitHub

←Índice de Artigos DPO→