Fonctionnalités de Trainer

Chaque recette ci-dessous illustre une fonctionnalité spécifique du Trainer : fonctions de perte personnalisées, évaluation économe en mémoire, stratégies de point de contrôle, et plus encore.

Ouvrez un issue si vous souhaitez voir une fonctionnalité ou un flux de travail spécifique ici.

Fonction de perte personnalisée

Passez compute_loss_func à Trainer pour remplacer la fonction de perte par défaut. Cette fonction s'exécute après le passage avant et ne définit que la manière dont la perte est calculée à partir des sorties. Pour modifier le passage avant lui-même, sous-classez plutôt compute_loss().

La fonction de perte personnalisée doit avoir la signature suivante :

import torch.nn.functional as F

def my_loss_fn(outputs, labels, num_items_in_batch):
    logits = outputs["logits"]
    loss = F.cross_entropy(logits, labels, reduction="sum")
    return loss / num_items_in_batch

outputs est la sortie brute du modèle (outputs.logits a la forme (batch, seq_len, vocab_size)).
labels correspond aux identifiants de jetons extraits du lot d’entrée par Trainer avant la passe avant.
num_items_in_batch est le nombre de cibles de prédiction sur l’ensemble du lot accumulé. Pour les modèles de LM causaux, il compte les étiquettes décalées (labels[..., 1:]), car le décalage des étiquettes laisse la position 0 de chaque séquence sans cible. Voir Échelle de la perte pour plus de détails. Trainer ignore la normalisation automatique de la perte lorsqu’une fonction de perte personnalisée est fournie, votre fonction doit donc gérer la normalisation directement.

trainer = Trainer(
    model=model,
    args=TrainingArguments(...),
    train_dataset=train_dataset,
    compute_loss_func=my_loss_fn,
)
trainer.train()

Consultez le guide sur le sous-classage pour plus d’exemples de surcharge de compute_loss().

Évaluation au démarrage

Définissez eval_on_start=True pour exécuter un passage d’évaluation complet avant la première étape d’entraînement. Une évaluation préalable à l’entraînement permet de détecter tôt les problèmes du pipeline d’évaluation, en particulier lors des exécutions de longue durée.

eval_on_start nécessite un eval_strategy valide (tel que "epoch") et un jeu de données d'évaluation.

from transformers import Trainer, TrainingArguments

trainer = Trainer(
    model=model,
    args=TrainingArguments(
        output_dir="out",
        eval_strategy="epoch",
        eval_on_start=True,
    ),
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    compute_metrics=compute_metrics,
)
trainer.train()

Un eval complet ajoute du temps, il est donc le plus utile lors des premières exécutions ou après modification de compute_metrics.

Évaluations économes en mémoire

Lors de l'évaluation, Trainer effectue un passage avant sur chaque lot et concatène les logits en un unique tenseur sur le GPU. Une fois le jeu de données d'évaluation entièrement traité, Trainer transfère les logits concaténés vers le CPU et appelle compute_metrics. Pour les grands modèles ou les grands jeux de données d'évaluation, les logits accumulés peuvent épuiser la mémoire GPU, même si l'entraînement sur le même matériel fonctionne correctement, car l'entraînement ne conserve qu'un seul lot d'activations à la fois.

eval_accumulation_steps

Déchargez les prédictions accumulées de la GPU vers la CPU toutes les n lots. Des valeurs plus faibles réduisent la mémoire GPU au prix de transferts CPU plus fréquents.

from transformers import Trainer, TrainingArguments

trainer = Trainer(
    model=model,
    args=TrainingArguments(
        output_dir="out",
        eval_strategy="epoch",
        eval_accumulation_steps=16,   # move predictions to CPU every 16 batches
    ),
    eval_dataset=eval_dataset,
    compute_metrics=compute_metrics,
)
trainer.train()

preprocess_logits_for_metrics

Appelé une fois par lot d'évaluation sur le GPU, immédiatement après le passage avant et avant l'accumulation des logits. La valeur retournée remplace les logits dans eval_pred.predictions. L'exécution du calcul au niveau du lot réduit la taille des tenseurs par lot et fournit à eval_accumulation_steps un tenseur plus petit à décharger.

import evaluate
from transformers import Trainer, TrainingArguments

metric = evaluate.load("accuracy")

def preprocess_logits_for_metrics(logits, labels):
    if isinstance(logits, tuple):
        logits = logits[0]
    return logits.argmax(dim=-1)

def compute_metrics(eval_preds):
    preds, labels = eval_preds
    labels = labels[:, 1:].reshape(-1)
    preds = preds[:, :-1].reshape(-1)
    return metric.compute(predictions=preds, references=labels)

trainer = Trainer(
    model=model,
    args=TrainingArguments(
        output_dir="out",
        eval_strategy="epoch",
        eval_accumulation_steps=16,
    ),
    eval_dataset=eval_dataset,
    compute_metrics=compute_metrics,
    preprocess_logits_for_metrics=preprocess_logits_for_metrics,
)
trainer.train()

Performance du Dataloader

Par défaut, Trainer crée un dataloader avec dataloader_num_workers=0. Les données sont chargées dans le processus principal pendant que le GPU reste inactif, ce qui se traduit par une faible utilisation du GPU entre les lots.

Les deux paramètres dataloader_persistent_workers et dataloader_prefetch_factor exigent que dataloader_num_workers > 0.

dataloader_persistent_workers maintient les sous-processus de travail en vie entre les époques afin d’éviter de les réinitialiser à chaque fois, au prix d’une consommation mémoire plus élevée.
dataloader_prefetch_factor contrôle le nombre de lots que chaque worker prépare à l’avance. Avec dataloader_prefetch_factor=2 et num_workers=4, jusqu’à 8 lots restent en mémoire pendant que le GPU s’entraîne sur le lot courant.

from transformers import TrainingArguments

args = TrainingArguments(
    output_dir="out",
    dataloader_num_workers=4,            # spawn 4 worker subprocesses
    dataloader_persistent_workers=True,  # keep them alive between epochs
    dataloader_prefetch_factor=2,        # each worker preloads 2 batches ahead
)

Grouper les échantillons par longueur

Utilisez train_sampling_strategy="group_by_length" pour regrouper par lots des exemples de longueurs similaires et réduire le bourrage. Lorsque vous ne fournissez pas de longueurs précalculées, Trainer les déduit à partir de la première entrée du modèle de chaque élément du jeu de données. Cela fonctionne également lorsque les jeux de données multimodaux basés sur un processeur renvoient des objets BatchFeature, car ce sont des conteneurs de caractéristiques de type mapping.

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="qwen3-vl-finetuned",
    train_sampling_strategy="group_by_length",
)

Si un Dataset possède déjà une colonne de longueur précalculée, Trainer utilise cette colonne à la place. Le nom de colonne par défaut est length. Définissez length_column_name si votre jeu de données utilise un autre nom. Cette stratégie nécessite un jeu de données dont la longueur est connue et est ignorée pour IterableDataset.

Échantillonnage de rééquilibrage par lot

Sur des jeux de données de longueurs variables, le déséquilibre au sein d’un micro-lot et entre les appareils entraîne une perte de temps des appareils sur le bourrage et une inactivité aux étapes de synchronisation des gradients.

Définissez train_sampling_strategy="batch_rebalance" dans TrainingArguments pour atténuer les deux effets. À chaque étape de l’optimiseur, le sampleur :

Trie les échantillons du lot par longueur.
Répartit le lot trié sur les différents dispositifs afin d’équilibrer le coût des jetons de bourrage pour chaque micro-lot. Les micro-lots contenant des échantillons longs reçoivent moins d’échantillons, tandis que ceux contenant des échantillons courts en reçoivent davantage.

Cela réduit le bourrage au sein de chaque micro-lot, et chaque appareil termine un micro-lot à peu près au même moment, ce qui réduit le temps d’attente lors de la synchronisation et maintient la mémoire de pointe plus basse que "group_by_length". Cette stratégie n’est pour l’instant prise en charge que pour l’entraînement parallèle par données (le parallélisme par tenseurs n’est pas encore pris en charge).

from transformers import Trainer, TrainingArguments

trainer = Trainer(
    model=model,
    args=TrainingArguments(
        output_dir="out",
        train_sampling_strategy="batch_rebalance",  # balance padding cost across devices
        per_device_train_batch_size=8,              # average samples per micro-batch
        length_column_name="length",                # optional: dataset column with precomputed lengths
    ),
    train_dataset=train_dataset,
)
trainer.train()

per_device_train_batch_size est ici une moyenne plutôt qu'un nombre exact par étape : certains micro-lots contiennent moins d'échantillons et d'autres en contiennent plus, mais le nombre total d'échantillons entraînés par étape reste identique à celui d'une stratégie d'échantillonnage distribué classique.

Le sampler a besoin de la longueur de chaque échantillon pour trier et équilibrer les lots. Par défaut, le Trainer parcourt l’ensemble de données complet une fois au début de l’entraînement pour calculer ces longueurs. Pour éviter ce parcours, précalculez les longueurs dans une colonne de l’ensemble de données (par exemple, lors de la prétraitement) et passez son nom en tant que length_column_name ("length" par défaut).

NEFTune

NEFTune ajoute du bruit aléatoire aux embeddings de jetons lors de la passe avant. Ce bruit agit comme une régularisation et peut améliorer les performances lors de l’ajustement fin sur des instructions.

Activez NEFTune en définissant neftune_noise_alpha dans TrainingArguments. Les valeurs typiques d’alpha se situent entre 5 et 15.

from transformers import Trainer, TrainingArguments

trainer = Trainer(
    model=model,
    args=TrainingArguments(
        output_dir="out",
        num_train_epochs=3,
        neftune_noise_alpha=5,
    ),
    train_dataset=train_dataset,
)
trainer.train()

NEFTune n’affecte que l’entraînement, et la couche d’embedding d’origine est restaurée après l’entraînement.

Journalisation

Contrôlez quand et où Trainer écrit les entrées de journalisation à l’aide de logging_strategy, logging_steps et report_to.

logging_strategy="steps" enregistre les journaux à chaque logging_steps() mises à jour de l’optimiseur. Utilisez "epoch" pour enregistrer les journaux à la fin de chaque époque à la place.
report_to transmet les journaux à un traceur d'expériences tel que Trackio.

from transformers import Trainer, TrainingArguments

trainer = Trainer(
    model=model,
    args=TrainingArguments(
        output_dir="out",
        logging_strategy="steps",
        logging_steps=50,               # write a log entry every 50 optimizer updates
        report_to="trackio",            # stream to Trackio (or "wandb", "tensorboard", …)
        run_name="model-experiment-v1", # display name in the tracker
    ),
    train_dataset=train_dataset,
)
trainer.train()

Checkpointing

Trainer enregistre un point de contrôle toutes les save_steps() mises à jour de l’optimiseur et les conserve tous (ou les ~TrainingArguments.save_total_limit plus récents).

save_strategy="best" ne conserve que le meilleur point de contrôle selon une métrique. Un nouveau point de contrôle n'est enregistré que lorsque la métrique suivie s'améliore, ce qui économise de l'espace disque et évite l'accumulation de points de contrôle obsolètes.

from transformers import Trainer, TrainingArguments

trainer = Trainer(
    model=model,
    args=TrainingArguments(
        output_dir="out",
        eval_strategy="epoch",
        save_strategy="best",
        metric_for_best_model="perplexity",   # save when eval perplexity improves
        greater_is_better=False,              # lower perplexity is better
        load_best_model_at_end=True,          # load the best weights after training finishes
    ),
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    compute_metrics=compute_metrics,          # must return {"perplexity": ...}
)
trainer.train()

Reprendre l'entraînement

Passez resume_from_checkpoint=True à train() si l’entraînement a été interrompu et que vous souhaitez le reprendre sans perdre les progrès réalisés. L’entraînement reprendra depuis le dernier point de contrôle présent dans output_dir.

trainer.train(resume_from_checkpoint=True)

Spécifiez un chemin de point de contrôle pour reprendre à partir d'un point particulier.

trainer.train(resume_from_checkpoint="out/checkpoint-1000")

Lors de la reprise, Trainer restaure l'état de l'optimiseur, l'état du planificateur et l'état du générateur aléatoire.

La reprise d’un point de contrôle nécessite les fichiers d’état de l’optimiseur et du planificateur dans le répertoire du point de contrôle. Si ces fichiers sont absents (par exemple, lorsque save_only_model=True), l’optimiseur repart de zéro.

Point de contrôle JIT

Avec un point de contrôle périodique (save_strategy="steps" ou "epoch"), vous perdez toute la progression de l’entraînement entre le dernier point de contrôle enregistré et l’interruption. Sur les clusters partagés avec des charges de travail préemptibles telles que Kueue, les tâches peuvent être interrompues à tout moment, ce qui signifie que cet intervalle peut représenter des heures de calcul gaspillées.

Le checkpointing JIT (Just-In-Time) comble cet écart. Lorsque le formateur reçoit un signal SIGTERM, il enregistre un checkpoint au point exact où l’entraînement a été interrompu, afin de reprendre avec une perte de progression minimale. Il fonctionne en complément du checkpointing périodique. Les sauvegardes périodiques protègent contre les pannes et les défaillances matérielles, tandis que les sauvegardes JIT protègent contre la préemption et les arrêts propres.

Activez-le en définissant enable_jit_checkpoint=True dans TrainingArguments.

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="your-model",
    enable_jit_checkpoint=True,
)

Lors de la réception de SIGTERM, Trainer attend la fin de l'étape d'entraînement en cours, enregistre un point de contrôle, puis arrête l'entraînement de manière propre. Un fichier sentinelle (checkpoint-is-incomplete.txt) est créé au début de l'enregistrement et supprimé une fois le point de contrôle entièrement écrit. Si un répertoire de point de contrôle contient encore ce fichier, l'enregistrement a été interrompu avant sa complétion. Trainer ne le vérifie pas automatiquement ; il est donc recommandé de l'inspecter vous-même avant de reprendre l'entraînement.

Reprendre depuis le point de contrôle JIT de la même manière que pour tout autre point de contrôle.

trainer.train(resume_from_checkpoint=True)

Vous devez configurer votre orchestrateur afin de laisser suffisamment de temps à la sauvegarde pour s'achever. La période d'arrêt propre par défaut de Kubernetes n'est que de 30 secondes, ce qui est généralement insuffisant pour les modèles de grande taille.

Kubernetes

Slurm

Calculez la période de grâce requise en additionnant la durée maximale possible d'une étape d'entraînement, le temps de sauvegarde du point de contrôle et le délai de 3 secondes de kill_wait avant le début de la sauvegarde. Par exemple, si une étape d'entraînement prend jusqu'à 2 minutes et que la sauvegarde d'un point de contrôle prend 2 minutes, définissez au moins 243 secondes de temps de grâce.

Update on GitHub

Transformers

Fonctionnalités de Trainer

Fonction de perte personnalisée

Évaluation au démarrage

Évaluations économes en mémoire

eval_accumulation_steps

preprocess_logits_for_metrics

Performance du Dataloader

Grouper les échantillons par longueur

Échantillonnage de rééquilibrage par lot

NEFTune

Journalisation

Checkpointing

Reprendre l'entraînement

Point de contrôle JIT