Funciones de Trainer

Cada receta a continuación demuestra una función específica de Trainer: funciones de pérdida personalizadas, evaluación eficiente en memoria, estrategias de checkpointing y más.

Abre un issue si hay una función o flujo de trabajo que te gustaría ver aquí.

Función de pérdida personalizada

Pase compute_loss_func a Trainer para reemplazar la función de pérdida predeterminada. La función se ejecuta después de la pasada hacia adelante y solo define cómo se calcula la pérdida a partir de las salidas. Para modificar la propia pasada hacia adelante, subclase compute_loss() en su lugar.

La función de pérdida personalizada debe tener la siguiente firma:

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 es la salida cruda del modelo (outputs.logits tiene forma (batch, seq_len, vocab_size)).
labels son los ids de tokens extraídos del lote de entrada por Trainer antes de la pasada hacia adelante.
num_items_in_batch es el número de objetivos de predicción en el lote acumulado completo. Para modelos de LM causal, cuenta las etiquetas desplazadas (labels[..., 1:]), ya que el desplazamiento de etiquetas deja la posición 0 de cada secuencia sin objetivo. Consulte Escala de la pérdida para más detalles. Trainer omite la normalización automática de la pérdida cuando se proporciona una función de pérdida personalizada, por lo que su función debe manejar la normalización directamente.

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

Consulte la guía de subclases para más ejemplos de cómo sobrescribir compute_loss().

Evaluación al inicio

Establezca eval_on_start=True para ejecutar una pasada completa de evaluación antes del primer paso de entrenamiento. Una evaluación previa al entrenamiento permite detectar problemas en el pipeline de evaluación de forma temprana, especialmente durante ejecuciones prolongadas.

eval_on_start requiere un eval_strategy válido (como "epoch") y un conjunto de datos de evaluación.

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

Una evaluación completa añade tiempo, por lo que resulta más útil en las primeras ejecuciones o después de modificar compute_metrics.

Evaluaciones eficientes en memoria

Durante la evaluación, Trainer ejecuta una pasada hacia adelante en cada lote y concatena los logits en un único tensor en la GPU. Una vez que el conjunto de datos de evaluación se ha procesado por completo, Trainer transfiere los logits concatenados a la CPU y llama a compute_metrics. Para modelos grandes o conjuntos de evaluación extensos, los logits acumulados pueden agotar la memoria de la GPU incluso cuando el entrenamiento en el mismo hardware funciona correctamente, ya que el entrenamiento solo mantiene un lote de activaciones a la vez.

eval_accumulation_steps

Descargue las predicciones acumuladas de la GPU a la CPU cada n lotes. Los valores más bajos reducen la memoria de la GPU a costa de transferencias a la CPU más frecuentes.

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

Se llama una vez por lote de evaluación en la GPU, inmediatamente después de la pasada hacia adelante y antes de la acumulación de logits. El valor devuelto reemplaza los logits en eval_pred.predictions. Ejecutar el cálculo a nivel de lote reduce el tamaño de los tensores por lote y proporciona a eval_accumulation_steps un tensor más pequeño para descargar.

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

Rendimiento del Dataloader

Por defecto, Trainer crea un dataloader con dataloader_num_workers=0. Los datos se cargan en el proceso principal mientras la GPU permanece inactiva, lo que se manifiesta como un bajo uso de la GPU entre lotes.

Tanto dataloader_persistent_workers como dataloader_prefetch_factor requieren dataloader_num_workers > 0.

dataloader_persistent_workers mantiene los subprocesos de los workers activos entre épocas para evitar reinicializarlos desde cero, a cambio de un mayor consumo de memoria.
dataloader_prefetch_factor controla cuántos lotes prepara cada worker por adelantado. Con dataloader_prefetch_factor=2 y num_workers=4, hasta 8 lotes permanecen en memoria mientras la GPU entrena con el lote actual.

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
)

Agrupar muestras por longitud

Use train_sampling_strategy="group_by_length" para agrupar en lotes ejemplos de longitudes similares y reducir el padding. Cuando no se proporcionan longitudes precalculadas, Trainer las infiere a partir de la primera entrada del modelo en cada elemento del conjunto de datos. Esto también funciona cuando los conjuntos de datos multimodales basados en procesadores devuelven objetos BatchFeature, ya que son contenedores de características con comportamiento de mapeo.

from transformers import TrainingArguments

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

Si un Dataset ya tiene una columna de longitud precalculada, Trainer utiliza esa columna en su lugar. El nombre de columna predeterminado es length. Establezca length_column_name si su conjunto de datos utiliza otro nombre. Esta estrategia requiere un conjunto de datos con una longitud conocida y se ignora para IterableDataset.

Muestreo de reequilibrio por lotes

En conjuntos de datos de longitud variable, el desequilibrio dentro de un micro-lote y entre dispositivos provoca que los dispositivos pierdan tiempo en el relleno (padding) y queden inactivos en las etapas de sincronización de gradientes.

Establezca train_sampling_strategy="batch_rebalance" en TrainingArguments para reducir ambos efectos. En cada paso del optimizador, el muestreador:

Ordena las muestras del lote por longitud.
Distribuye el lote ordenado entre los dispositivos de modo que el coste de los tokens rellenos de cada micro-lote quede equilibrado. Los micro-lotes con muestras largas reciben menos muestras, y los micro-lotes con muestras cortas reciben más.

Esto reduce el relleno dentro de cada micro-lote y cada dispositivo termina un micro-lote aproximadamente al mismo tiempo, lo que reduce el tiempo de inactividad en la sincronización y mantiene la memoria pico más baja que "group_by_length". Esta estrategia solo se admite actualmente para el entrenamiento paralelo de datos (la paralelización de tensores aún no se admite).

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 es un promedio aquí, no un recuento exacto por paso: algunos micro-lotes tienen menos muestras y otros tienen más, pero el número total de muestras entrenadas por paso se mantiene igual que con una estrategia de muestreo distribuido normal.

El sampler necesita la longitud de cada muestra para ordenar y equilibrar los lotes. De forma predeterminada, el Trainer recorre el conjunto de datos completo una vez al inicio del entrenamiento para calcular estas longitudes. Para omitir este recorrido, precalcula las longitudes en una columna del conjunto de datos (por ejemplo, durante el preprocesamiento) y pasa su nombre como length_column_name ("length" de forma predeterminada).

NEFTune

NEFTune añade ruido aleatorio a las incrustaciones de tokens durante la pasada hacia adelante. El ruido actúa como regularización y puede mejorar el rendimiento en el ajuste fino de instrucciones.

Active NEFTune configurando neftune_noise_alpha en TrainingArguments. Los valores típicos de alpha van de 5 a 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 solo afecta al entrenamiento y la capa de incrustación original se restaura después del entrenamiento.

Registro

Controle cuándo y dónde Trainer escribe las entradas de registro con logging_strategy, logging_steps y report_to.

logging_strategy="steps" registra cada logging_steps() actualizaciones del optimizador. Utilice "epoch" para registrar al final de cada época en su lugar.
report_to envía los registros a un rastreador de experimentos como 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()

Puntos de control

Trainer guarda un punto de control cada save_steps() actualizaciones del optimizador y conserva todos ellos (o los más recientes ~TrainingArguments.save_total_limit).

save_strategy="best" conserva únicamente el mejor punto de control según una métrica. Se guarda un nuevo punto de control solo cuando la métrica de seguimiento mejora, lo que ahorra espacio en disco y evita la acumulación de puntos de control obsoletos.

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

Reanudar el entrenamiento

Pasa resume_from_checkpoint=True a train() si el entrenamiento se interrumpió y deseas reanudarlo sin perder el progreso. El entrenamiento se reanudará desde el último punto de control en output_dir.

trainer.train(resume_from_checkpoint=True)

Especifique una ruta de punto de control para reanudar desde un punto determinado.

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

Al reanudar, Trainer restaura el estado del optimizador, el estado del programador y el estado del RNG.

Reanudar desde un checkpoint requiere los archivos de estado del optimizador y del programador en el directorio del checkpoint. Si esos archivos no están presentes (por ejemplo, cuando save_only_model=True), el optimizador se reinicia desde cero.

Puntos de control JIT

Con el guardado periódico de puntos de control (save_strategy="steps" o "epoch"), se pierde todo el progreso de entrenamiento entre el último punto de control guardado y una interrupción. En clústeres compartidos con cargas de trabajo preemibles, como Kueue, los trabajos pueden terminarse en cualquier momento, por lo que esa brecha puede suponer horas de cómputo desperdiciado.

El checkpointing JIT (Just-In-Time) cierra esta brecha. Cuando el trainer recibe una señal SIGTERM, guarda un checkpoint en el punto exacto en que se interrumpió el entrenamiento, de modo que se reanuda con una pérdida mínima de progreso. Funciona junto con el checkpointing periódico. Las guardas periódicas protegen contra fallos y problemas de hardware, mientras que las guardas JIT protegen contra la preemption y los apagados controlados.

Actívalo estableciendo enable_jit_checkpoint=True en TrainingArguments.

from transformers import TrainingArguments

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

Cuando se recibe SIGTERM, Trainer espera a que el paso de entrenamiento actual termine, guarda un punto de control y detiene el entrenamiento de forma ordenada. Se escribe un archivo centinela (checkpoint-is-incomplete.txt) al comenzar el guardado y se elimina una vez que el punto de control se ha escrito por completo. Si un directorio de punto de control aún contiene este archivo, el guardado se interrumpió antes de completarse. Trainer no lo comprueba automáticamente, por lo que debe inspeccionarlo usted mismo antes de reanudar.

Reanuda desde el punto de control de JIT de la misma manera que cualquier otro punto de control.

trainer.train(resume_from_checkpoint=True)

Debes configurar tu orquestador para permitir suficiente tiempo para que el checkpoint se complete. El período de apagado ordenado por defecto en Kubernetes es de solo 30 segundos, lo cual generalmente no es suficiente para modelos más grandes.

Kubernetes

Slurm

Calcule el período de gracia necesario como la suma del tiempo máximo posible de un paso de entrenamiento, el tiempo de guardado del punto de control y el retardo de 3 segundos de kill_wait antes de que comience el guardado. Por ejemplo, si un paso de entrenamiento tarda hasta 2 minutos y guardar un punto de control tarda 2 minutos, establezca al menos 243 segundos de tiempo de gracia.

Update on GitHub

Transformers

Funciones de Trainer

Función de pérdida personalizada

Evaluación al inicio

Evaluaciones eficientes en memoria

eval_accumulation_steps

preprocess_logits_for_metrics

Rendimiento del Dataloader

Agrupar muestras por longitud

Muestreo de reequilibrio por lotes

NEFTune

Registro

Puntos de control

Reanudar el entrenamiento

Puntos de control JIT