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

Aperçu

Le Distillation Trainer implémente la distillation de connaissances on-policy telle que décrite dans On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes par Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos, Matthieu Geist et Olivier Bachem.

L'abstract de l'article est le suivant :

La distillation de connaissances (KD) est largement utilisée pour compresser un modèle enseignant afin de réduire son coût d’inférence et son empreinte mémoire, en entraînant un modèle étudiant plus petit. Cependant, les méthodes actuelles de KD pour les modèles de séquences auto-régressifs souffrent d’un décalage de distribution entre les séquences de sortie observées pendant l’entraînement et celles générées par l’étudiant lors de l’inférence. Pour remédier à ce problème, nous introduisons la distillation de connaissances généralisée (GKD). Plutôt que de s’appuyer uniquement sur un ensemble fixe de séquences de sortie, GKD entraîne l’étudiant sur ses propres séquences de sortie auto-générées en exploitant les retours de l’enseignant sur ces séquences. Contrairement aux approches de KD supervisées, GKD offre également la flexibilité d’employer des fonctions de perte alternatives entre l’étudiant et l’enseignant, ce qui peut s’avérer utile lorsque l’étudiant manque d’expressivité pour imiter la distribution de l’enseignant.

Le DistillationTrainer entraîne un modèle étudiant plus petit afin qu’il reproduise la distribution des jetons suivants du modèle enseignant sur les générations on-policy de l’étudiant. Il génère les complétions de l’étudiant on-policy (facultativement propulsées par vLLM) et aligne l’enseignant sur sa distribution complète des jetons suivants à l’aide d’une perte de divergence de Jensen-Shannon par morceaux, économe en mémoire, de sorte que la distribution dense de l’enseignant n’est jamais matérialisée en totalité.

Ce formateur a été contribué par Carlos Miguel Patiño.

Pour une variante asynchrone qui découple la génération de l’entraînement et évalue les sorties par rapport à un enseignant servi via HTTP plutôt qu’à un passage avant local, consultez AsyncDistillationTrainer.

Démarrage rapide

Cet exemple montre comment entraîner un modèle à l'aide de la méthode de distillation. Nous distillons un modèle Qwen 2.5 0.5B Instruct à partir d'un enseignant Qwen 2.5 1.5B Instruct sur les prompts du jeu de données de prompts UltraFeedback. Vous pouvez consulter les données du jeu de données ici :

Voici le script pour entraîner le modèle.

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

Exécutez le script à l'aide de la commande suivante :

accelerate launch train_distillation.py

Explorer plus en profondeur la méthode de distillation

La distillation de connaissances on-policy entraîne un modèle élève à reproduire la distribution des jetons suivants d’un modèle enseignant sur des complétions générées par l’élève lui-même, plutôt que sur un ensemble fixe de sorties de l’enseignant. Apprendre à partir de ses propres générations permet à l’élève de corriger ses propres erreurs, ce qui surpasse généralement la distillation off-policy. Cette section détaille le fonctionnement en pratique, en couvrant les deux étapes clés : la génération des complétions et le calcul de la perte.

Générer des complétions

À chaque étape d'entraînement, l'élève génère un lot de complétions pour les prompts échantillonnés.

Calcul de la perte

La perte est la divergence de Jensen-Shannon généralisée (JSD) entre la distribution de l'élève $p_S$ et la distribution de l'enseignant $p_T$ sur les jetons de complétion générés, interpolée par beta et définie comme suit : $\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,$

où $p_M$ est le $\beta$ -mélange des deux distributions. Les extrémités se réduisent aux divergences pures : beta=0.0 donne la KL avant $\mathbb{D}_{\mathrm{KL}}\!\left[ p_T \| p_S \right]$ et beta=1.0 la KL inverse $\mathbb{D}_{\mathrm{KL}}\!\left[ p_S \| p_T \right]$ .

En pratique, la projection vers les logits du vocabulaire et la divergence sont calculés par lots, de sorte que la mémoire d'activation de pointe n'évolue pas avec le tenseur complet de logits vocabulaire × longueur de séquence. Voir Réduire l'utilisation de la mémoire.

Type de jeu de données attendu

Le jeu de données doit être formaté comme un jeu de données conversationnel contenant uniquement des prompts. L’étudiant génère ses propres complétions on-policy, il n’est donc nécessaire que de fournir le prompt :

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

Métriques journalisées

Lors de l'entraînement et de l'évaluation, nous enregistrons les métriques suivantes :

num_tokens : le nombre total de jetons traités jusqu'à présent, incluant à la fois les invites et les complétions. Lors de l'utilisation d'outils, seuls les jetons non liés aux outils sont comptés.
step_time : le temps moyen (en secondes) par étape d'entraînement (génération incluse).
completions/mean_length : La longueur moyenne des complétions générées. Lors de l'utilisation d'outils, seuls les jetons non-outils sont comptés.
completions/min_length : La longueur minimale des complétions générées. Lors de l'utilisation d'outils, seuls les jetons non-outils sont comptés.
completions/max_length : La longueur maximale des complétions générées. Lors de l'utilisation d'outils, seuls les jetons non-outils sont comptés.
completions/mean_terminated_length : La longueur moyenne des complétions générées qui se terminent par EOS. Lors de l'utilisation d'outils, seuls les jetons non-outils sont comptés.
completions/min_terminated_length : La longueur minimale des complétions générées qui se terminent par EOS. Lors de l'utilisation d'outils, seuls les jetons non-outils sont comptés.
completions/max_terminated_length : La longueur maximale des complétions générées qui se terminent par EOS. Lors de l'utilisation d'outils, seuls les jetons non-outils sont comptés.
completions/clipped_ratio : Le ratio des complétions tronquées (clippées).
tools/call_frequency : nombre moyen d'appels d'outils par complétion dans le lot de génération. Journalisé uniquement lorsque des tools sont fournis.
tools/failure_frequency : La proportion d'appels d'outils qui ont échoué (l'outil n'a pas été trouvé, a levé une exception, ou le type d'appel n'est pas pris en charge). Elle est 0.0 lorsqu'aucun outil n'a été appelé. Journalisé uniquement lorsque des tools sont fournis.
entropy : Entropie moyenne des prédictions de jetons sur les complétions générées (en nats).

Personnalisation

Accélérer l'entraînement avec la génération propulsée par vLLM

La génération est souvent le principal goulot d'étranglement lors de l'entraînement avec des méthodes on-policy. Pour accélérer la génération, vous pouvez utiliser vLLM, un moteur d'inférence à haut débit et faible latence pour les LLM. Pour l'activer, installez d'abord le package avec

pip install trl[vllm]

Nous prenons en charge deux façons d'utiliser vLLM pendant l'entraînement : le mode colocalisé et le mode serveur.

Option 1 : Mode Colocate

Dans ce mode, vLLM s’exécute dans le processus d’entraînement et partage la mémoire GPU avec le modèle d’entraînement. Cela évite de démarrer un serveur séparé et peut améliorer l’utilisation des GPU, mais peut entraîner une concurrence mémoire sur les GPU d’entraînement. C’est le mode par défaut.

from trl import DistillationConfig

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

Option 2 : Mode serveur

Dans ce mode, vLLM s’exécute dans un processus distinct (et utilise des GPU séparés) et communique avec l’entraîneur via HTTP. C’est idéal si vous disposez de GPU dédiés à l’inférence.

Démarrer le serveur vLLM :

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

Activer le mode serveur dans votre script d’entraînement :

from trl import DistillationConfig

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

Assurez-vous que le serveur utilise des GPU différents de ceux du formateur, sinon vous risquez de rencontrer des erreurs NCCL. Vous pouvez spécifier les GPU à utiliser avec la variable d’environnement CUDA_VISIBLE_DEVICES.

En fonction de la taille du modèle et des besoins globaux en mémoire GPU pour l’entraînement, vous devrez peut-être ajuster le paramètre vllm_gpu_memory_utilization dans DistillationConfig afin d’éviter une sous-utilisation ou des erreurs de mémoire insuffisante.

Pour plus d'informations, consultez Accélérer l'entraînement avec vLLM.

Entraîner des adaptateurs avec PEFT

Nous prenons en charge une intégration étroite avec la bibliothèque 🤗 PEFT, vous permettant d’entraîner des adaptateurs et de les partager sur le Hub plutôt que d’entraîner l’ensemble de l’étudiant.

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

La perte de distillation lit directement lm_head.weight et exécute le tronc de l'étudiant sans passer par PeftModel.forward(). Les adaptateurs sur lm_head (via target_modules) et les méthodes d'apprentissage de prompt (PromptTuning, PrefixTuning, P-Tuning) sont donc rejetés, car ils seraient ignorés silencieusement. Pour entraîner la tête, utilisez plutôt modules_to_save=["lm_head"].

Entraînement des agents

DistillationTrainer prend en charge l'entraînement d'agents : l'élève appelle des outils pendant la génération et est distillé sur la trajectoire résultante. Les jetons de résultat d'outil sont masqués dans la perte, de sorte que l'élève n'est entraîné que sur les jetons qu'il a lui-même générés.

Outils

L'argument tools attend une liste de fonctions Python qui définissent les outils disponibles pour l'agent :

from trl import DistillationTrainer

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

Chaque outil doit être une fonction Python standard avec des arguments et types de retour annotés, accompagnée d'une docstring au style Google décrivant son objectif, ses arguments et sa valeur de retour. Pour plus de détails, consultez le guide sur la transmission d'outils.

La boucle d’appel d’outils exige que le modèle de discussion soit préfixe-préservant (l’ajout d’un message d’outil ne doit pas modifier la façon dont les messages précédents sont rendus). Pour les familles de modèles connues (par ex. Qwen3, DeepSeek-V3), TRL remplace automatiquement le modèle d’entraînement par un modèle corrigé lorsque les outils sont activés. Consultez Chat Templates pour la liste complète.

Utilisez max_tool_calling_iterations dans DistillationConfig pour limiter le nombre de tours d’appel d’outils. Par défaut, il n’y a pas de limite et la génération s’arrête lorsque l’étudiant produit un tour de réponse sans appel d’outils.

Exemple :

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],
    ...,
)

Les outils asynchrones ne sont pas encore pris en charge — passez des fonctions synchrones.

Réponses d'outils multimodales

Les outils peuvent renvoyer des images en plus du texte en retournant une liste de blocs de contenu. Cela est utile pour l’entraînement d’agents VLM où l’outil fournit un retour visuel (par exemple, des captures d’écran, des graphiques, des images de caméra).

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."}]

Les images renvoyées sont automatiquement injectées dans la conversation et transmises au VLM pour les tours de génération suivants.

Entraînement des modèles de vision et de langage

DistillationTrainer prend en charge la distillation de modèles Vision-Language (VLM) sur des jeux de données multimodaux contenant à la fois du texte et des images. Transmettez un VLM à la fois comme élève et comme enseignant, et fournissez un jeu de données ne contenant que des prompts, avec soit une colonne image (une image par échantillon), soit une colonne images (liste d’images par échantillon). Pour plus d’informations sur la structure attendue du jeu de données, consultez la section Format du jeu de données — Jeux de données Vision.

Testé avec :

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

La compatibilité avec tous les VLM n'est pas garantie. Si vous pensez qu'un modèle devrait être pris en charge, n'hésitez pas à ouvrir un ticket sur GitHub — ou mieux encore, à soumettre une pull request avec les modifications nécessaires.

Interface en ligne de commande

Utilisez l'interface en ligne de commande trl distillation pour lancer l'entraînement par distillation depuis le terminal. Elle prend en charge l'entraînement complet et LoRA via les drapeaux standard 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 )

Paramètres

model (str ou PreTrainedModel ou PeftModel) — Modèle à entraîner. Peut être l'un des éléments suivants :
- Une chaîne de caractères, correspondant à l'identifiant du modèle d'un modèle pré-entraîné hébergé dans un dépôt de modèles sur huggingface.co, ou à un chemin vers un dossier contenant les poids du modèle enregistrés à l'aide de save_pretrained, par exemple './my_model_directory/'. Le modèle est chargé à l'aide de <ModelArchitecture>.from_pretrained (où <ModelArchitecture> est dérivé de la configuration du modèle) avec les arguments nommés de args.model_init_kwargs. Si dtype n'est pas spécifié dans args.model_init_kwargs, la valeur par défaut est float32. Cela diffère de from_pretrained, où (depuis Transformers v5) le type de données est déduit de la configuration du modèle.
- Un objet PreTrainedModel. Seuls les modèles de langage causaux sont pris en charge.
- Un objet PeftModel. Seuls les modèles de langage causaux sont pris en charge.
teacher_model (str ou PreTrainedModel, facultatif) — Modèle enseignant dont la distribution des jetons suivants l’élève est entraîné à imiter. Peut être un identifiant de modèle / chemin (chargé comme model, en utilisant args.teacher_model_init_kwargs) ou un PreTrainedModel instancié. Il doit partager le vocabulaire de l’élève. Peut être omis par les sous-classes qui fournissent l’enseignant autrement (par ex. un serveur distant).
args (DistillationConfig, facultatif) — Configuration de ce formateur. Si None, une configuration par défaut est utilisée.
train_dataset (Dataset ou IterableDataset) — Jeu de données à utiliser pour l'entraînement. Il doit inclure une colonne "prompt". Toute colonne supplémentaire dans le jeu de données est ignorée. Le format des échantillons peut être l'un des deux suivants :
- Standard : chaque échantillon contient du texte brut.
- Conversationnel : chaque échantillon contient des messages structurés (par exemple, le rôle et le contenu).
Lorsque train_dataset est un IterableDataset (par exemple un jeu de données en streaming), max_steps doit être défini dans les arguments d’entraînement, car sa longueur ne peut pas être déduite et le nombre total d’étapes d’entraînement est nécessaire pour limiter la boucle d’entraînement et configurer le planificateur de taux d’apprentissage.
eval_dataset (Dataset, IterableDataset, DatasetDict, IterableDatasetDict ou dict[str, Dataset | IterableDataset]) — Jeu de données à utiliser pour l'évaluation. Il doit répondre aux mêmes exigences que train_dataset.
processing_class (PreTrainedTokenizerBase, ProcessorMixin, facultatif) — Classe de traitement utilisée pour traiter les données. Le côté de bourrage doit être défini sur « left ». Si None, la classe de traitement est chargée à partir du nom du modèle avec from_pretrained. Un jeton de bourrage, tokenizer.pad_token, doit être défini. Si la classe de traitement n’a pas défini de jeton de bourrage, tokenizer.eos_token sera utilisé par défaut.
callbacks (liste de TrainerCallback, facultatif) — Liste de callbacks pour personnaliser la boucle d'entraînement. Ils seront ajoutés à la liste des callbacks par défaut détaillée ici.
Si vous souhaitez supprimer l'un des callbacks par défaut utilisés, utilisez la méthode remove_callback.
optimizers (tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None], facultatif, par défaut (None, None)) — Un tuple contenant l'optimiseur et le planificateur à utiliser. Par défaut, une instance de AdamW est utilisée sur votre modèle ainsi qu'un planificateur défini par get_linear_schedule_with_warmup et contrôlé par args.
quantization_config (BitsAndBytesConfig, facultatif) — Configuration de quantification utilisée lors du chargement du modèle à partir d'un identifiant de modèle. À combiner avec peft_config pour l'entraînement QLoRA. Ignoré si le modèle est déjà instancié.
peft_config (PeftConfig, facultatif) — Configuration PEFT utilisée pour envelopper le modèle. Si None, le modèle n'est pas enveloppé.
tools (liste de Callable, facultatif) — Une liste de fonctions d’outils appelables que le modèle peut invoquer pendant la génération. Chaque outil doit être une fonction Python standard avec des arguments et des valeurs de retour correctement annotés par des types, ainsi qu’une docstring au style Google décrivant son objectif, ses arguments et sa valeur de retour. Pour plus de détails, voir : https://huggingface.co/docs/transformers/en/chat_extras#passing-tools. Le modèle utilise le nom de la fonction, ses annotations de types et sa docstring pour déterminer comment l’appeler. Assurez-vous que le modèle de chat du modèle prend en charge l’utilisation d’outils et qu’il a été affiné pour l’appel d’outils.

Trainer pour la distillation de connaissances. L’élève est entraîné on-policy — il génère lui-même les complétions — afin de correspondre à la distribution du prochain jeton de l’enseignant sous une divergence de Jensen-Shannon généralisée (interpolant KL avant, KL arrière et JSD via beta), telle qu’introduite dans On-Policy Distillation of Language Models.

Exemple :

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

train

< source >

Paramètres

resume_from_checkpoint (str ou bool, facultatif) — Si c'est une str, chemin local d'un point de contrôle enregistré par une instance précédente de Trainer. Si c'est un bool et vaut True, charge le dernier point de contrôle dans args.output_dir enregistré par une instance précédente de Trainer. S'il est présent, l'entraînement reprendra à partir des états du modèle/optimiseur/planificateur chargés ici.
trial (optuna.Trial ou dict[str, Any], facultatif) — L’exécution de l’essai ou le dictionnaire d’hyperparamètres pour la recherche d’hyperparamètres.
ignore_keys_for_eval (list[str], facultatif) — Une liste de clés dans la sortie de votre modèle (s’il s’agit d’un dictionnaire) qui doivent être ignorées lors de la collecte des prédictions pour l’évaluation pendant l’entraînement.

Retourne

~trainer_utils.TrainOutput

Objet contenant le nombre d'étapes globales, la perte d'entraînement et les métriques.

Point d'entrée principal de l'entraînement.

save_model

< source >

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

Enregistrera le modèle, ce qui vous permettra de le recharger à l'aide de from_pretrained().

N'enregistrera que depuis le processus principal.

push_to_hub

< source >

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

Paramètres

commit_message (str, facultatif, par défaut "End of training") — Message à committer lors de la poussée.
blocking (bool, facultatif, par défaut True) — Indique si la fonction doit ne renvoyer que lorsque le git push est terminé.
token (str, facultatif, par défaut None) — Jeton avec permission d’écriture pour écraser les arguments originaux du Trainer.
revision (str, facultatif) — La révision git à partir de laquelle effectuer le commit. Par défaut, il s'agit de la tête de la branche « main ».
kwargs (dict[str, Any], facultatif) — Arguments supplémentaires passés à ~Trainer.create_model_card.

Téléchargez self.model et self.processing_class sur le hub de modèles 🤗 dans le dépôt 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 )

Paramètres qui contrôlent le modèle et le modèle enseignant

model_init_kwargs (str ou dict[str, Any], facultatif) — Arguments par mot-clé pour AutoModelForCausalLM.from_pretrained, utilisés lorsque l’argument model du formateur est fourni sous forme de chaîne de caractères. La valeur de revision est également utilisée lors du chargement de la classe de traitement.
trust_remote_code (bool, facultatif, valeur par défaut False) — Indique s'il faut autoriser le chargement de modèles et de tokenizers qui incluent du code Python personnalisé depuis le Hub. Transmis à from_pretrained et from_pretrained, pour l'étudiant comme pour l'enseignant.
teacher_model_name_or_path (str, facultatif) — Nom ou chemin du modèle enseignant. Utilisé lorsque l’enseignant est chargé localement.
teacher_model_revision (str, facultatif) — Révision du modèle enseignant (par exemple, nom de branche, étiquette ou hachage de commit).
teacher_model_init_kwargs (str ou dict[str, Any], facultatif) — Arguments par mot-clé transmis à AutoModelForCausalLM.from_pretrained lors de l'instantiation du modèle enseignant à partir d'une chaîne de caractères.
disable_dropout (bool, facultatif, valeur par défaut False) — Indique s'il faut désactiver le dropout dans le modèle étudiant pendant l'entraînement.

Paramètres qui contrôlent la prétraitement des données

remove_unused_columns (bool, facultatif, valeur par défaut False) — Indique s'il faut conserver uniquement la colonne "prompt" dans le jeu de données. Le formateur consomme la colonne de prompt brute et génère des complétions on-policy, la valeur par défaut est donc False.
max_completion_length (int ou None, facultatif, par défaut 512) — Nombre maximal de jetons à générer par complétion lors de la génération on-policy.
ds3_gather_for_generation (bool, facultatif, valeur par défaut True) — Ce paramètre s’applique à DeepSpeed ZeRO-3. S’il est activé, les poids du modèle de politique sont regroupés pour la génération, ce qui améliore la vitesse de génération. Toutefois, désactiver cette option permet d’entraîner des modèles dont la taille dépasse la capacité VRAM d’une seule GPU, au prix d’une génération plus lente. La désactivation de cette option n’est pas compatible avec la génération vLLM.
shuffle_dataset (bool, facultatif, par défaut True) — Indique s'il faut mélanger le jeu de données d'entraînement.
pad_to_multiple_of (int ou str, facultatif, valeur par défaut "auto") — Si défini, les identifiants des prompts et des complétions seront rembourrés jusqu'à un multiple de cette valeur. "auto" rembourde jusqu'à un multiple de 128 sur Neuron (AWS Trainium), qui compile une fois par forme d'entrée, et ne rembourde pas ailleurs.

Paramètres qui contrôlent la génération

temperature (float, facultatif, par défaut 1.0) — Température pour l'échantillonnage lors de la génération et pour le calcul de la perte de distillation. Des valeurs plus élevées produisent des distributions de probabilité plus douces.
top_p (float, facultatif, valeur par défaut 1.0) — Paramètre d'échantillonnage Top-p (noyau) pour la génération on-policy.
top_k (int, facultatif, valeur par défaut 0) — Paramètre d'échantillonnage top-k pour la génération on-policy. 0 désactive le filtrage top-k.
min_p (float, facultatif) — Probabilité minimale du jeton, qui sera mise à l'échelle par la probabilité du jeton le plus probable. Elle doit être une valeur comprise entre 0.0 et 1.0. Les valeurs typiques se situent dans la plage 0.01-0.2.
generation_kwargs (dict[str, Any], facultatif) — Arguments supplémentaires à passer à GenerationConfig (si vous utilisez transformers) ou SamplingParams (si vous utilisez vLLM) lors de l'échantillonnage des complétions. Cela permet de personnaliser davantage le comportement de la génération, par exemple en définissant suppress_tokens, num_beams, etc. S'il contient des clés qui entrent en conflit avec les autres paramètres de génération (comme min_p, top_p, etc.), ils les remplaceront.
chat_template_kwargs (dict[str, Any], facultatif) — Arguments supplémentaires à passer à la fonction apply_chat_template lors de la génération des complétions.
repetition_penalty (float, facultatif, par défaut 1.0) — Flottant qui pénalise les nouveaux jetons en fonction de leur présence dans le prompt et dans le texte généré jusqu’à présent. Les valeurs > 1.0 encouragent le modèle à utiliser de nouveaux jetons, tandis que les valeurs < 1.0 encouragent le modèle à répéter les jetons.
cache_implementation (str, facultatif) — Implémentation de la méthode de cache pour une génération plus rapide lorsque use_vllm est défini sur False.

Paramètres qui contrôlent l'accélération de la génération propulsée par vLLM

use_vllm (bool, facultatif, valeur par défaut False) — Indique s'il faut utiliser vLLM pour générer des complétions on-policy à partir du modèle étudiant.
vllm_mode (str, facultatif, par défaut "colocate") — Mode d'intégration de vLLM pour l'étudiant. Soit "server", soit "colocate".
vllm_model_impl (str, facultatif, par défaut "vllm") — Backend d’implémentation du modèle pour vLLM. Utilisez "vllm" ou "transformers".
vllm_enable_sleep_mode (bool, facultatif, valeur par défaut False) — Active le mode veille de vLLM pour décharger les poids de l’étudiant pendant l’étape d’optimiseur.
vllm_structured_outputs_regex (str, facultatif) — Motif d'expression régulière pour les sorties structurées de vLLM.

Paramètres qui contrôlent le serveur vLLM (utilisés uniquement lorsque `vllm_mode` est `"server"`)

vllm_server_base_url (str, facultatif) — URL de base du serveur vLLM de l’étudiant. Si fournie, vllm_server_host et vllm_server_port sont ignorés.
vllm_server_host (str, facultatif, par défaut "0.0.0.0") — Hôte du serveur vLLM de l’étudiant.
vllm_server_port (int, facultatif, par défaut 8000) — Port du serveur vLLM de l'étudiant.
vllm_server_timeout (float, facultatif, par défaut 240.0) — Délai d’attente pour la connexion au serveur vLLM de l’étudiant.
vllm_group_port (int, facultatif, par défaut 51216) — Port du groupe de mise à jour des poids vLLM (communicateur NCCL).

Paramètres qui contrôlent l'exécution de vLLM co-localisée (utilisés uniquement lorsque `vllm_mode` est `"colocate"`)

vllm_gpu_memory_utilization (float, facultatif, valeur par défaut 0.3) — Taux d'utilisation de la mémoire GPU pour le moteur vLLM étudiant co-localisé.
vllm_max_model_length (int, facultatif) — Longueur maximale de séquence du modèle pour le moteur vLLM co-localisé.
vllm_tensor_parallel_size (int, facultatif, valeur par défaut 1) — Taille du parallélisme tensoriel pour le moteur vLLM étudiant co-localisé.

Paramètres qui contrôlent l'entraînement

beta (float, facultatif, par défaut 1.0) — Coefficient d’interpolation pour la perte de divergence de Jensen-Shannon généralisée. Lorsque 0.0, la perte est la divergence KL avant. Lorsque 1.0, la perte est la divergence KL inverse. Lorsque 0.5, il s’agit de la JSD standard. Contrairement au beta de GRPO (un coefficient de pénalité KL par rapport à un modèle de référence), ici il sélectionne la divergence elle-même ; il n’y a pas de pénalité KL par rapport à un modèle de référence.
max_tool_calling_iterations (int, facultatif) — Nombre maximal de tours d’appel d’outils lors de l’entraînement d’un agent. Si None, il n’y a pas de limite et la génération s’arrête lorsque le modèle génère un tour de réponse sans appels d’outils ou lorsque la longueur totale de la réponse atteint max_model_length.

Paramètres qui contrôlent la journalisation

log_completions (bool, facultatif, par défaut False) — Indique s'il faut consigner un échantillon de paires (prompt, completion) toutes les logging_steps étapes. Si rich est installé, l'échantillon est affiché. Si la journalisation wandb et/ou trackio est activée, il est consigné dans wandb et/ou trackio.
num_completions_to_print (int, facultatif) — Nombre de complétions à afficher avec rich. Si None, toutes les complétions sont journalisées.
log_unique_prompts (bool, facultatif, valeur par défaut False) — Indique s'il faut journaliser les prompts uniques. Si True, seuls les prompts uniques sont journalisés. Si False, tous les prompts sont journalisés.

Classe de configuration pour le DistillationTrainer.

Étend TrainingArguments avec des paramètres spécifiques à la distillation de connaissances. Tous les champs nécessaires sont déclarés ici.

À l’aide de HfArgumentParser, nous pouvons transformer cette classe en arguments argparse qui peuvent être spécifiés sur la ligne de commande.

Ces paramètres ont des valeurs par défaut différentes de TrainingArguments :

logging_steps : la valeur par défaut est 10 au lieu de 500.

gradient_checkpointing : la valeur par défaut est True au lieu de False.

bf16 : passe à True par défaut si fp16 n’est pas défini, au lieu de False.

learning_rate : la valeur par défaut est 1e-6 au lieu de 5e-5.

Update on GitHub

←Index des articles DPO→