Recursos do Trainer

Cada receita abaixo demonstra um recurso específico do Trainer: funções de perda personalizadas, avaliação eficiente em memória, estratégias de checkpointing e muito mais.

Abra uma issue se houver um recurso ou fluxo de trabalho que você gostaria de ver aqui.

Função de perda personalizada

Passe compute_loss_func para Trainer para substituir a função de perda padrão. A função é executada após o forward pass e define apenas como a perda é calculada a partir das saídas. Para modificar o próprio forward pass, crie uma subclasse de compute_loss() em vez disso.

A função de perda personalizada deve ter a seguinte assinatura:

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 é a saída bruta do modelo (outputs.logits tem a forma (batch, seq_len, vocab_size)).
• labels são os ids de tokens removidos do lote de entrada pelo Trainer antes da passagem para frente.
• num_items_in_batch é o número de alvos de previsão em todo o lote acumulado. Para modelos de LM causal, ele conta os rótulos deslocados (labels[..., 1:]), pois o deslocamento dos rótulos deixa a posição 0 de cada sequência sem um alvo. Consulte Escala da perda para mais detalhes. O Trainer ignora a normalização automática da perda quando uma função de perda personalizada é fornecida, portanto, sua função deve tratar a normalização diretamente.

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

Consulte o guia de subclasseamento para mais exemplos de como sobrescrever compute_loss().

Avaliando no início

Defina eval_on_start=True para executar uma passagem completa de avaliação antes do primeiro passo de treinamento. Uma avaliação pré-treinamento revela problemas no pipeline de avaliação precocemente, especialmente durante execuções longas.

eval_on_start requer um eval_strategy válido (como "epoch") e um conjunto de dados de avaliação.

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

Uma avaliação completa adiciona tempo, portanto, é mais útil em execuções iniciais ou após modificar compute_metrics.

Avaliações com uso eficiente de memória

Durante a avaliação, o Trainer executa uma passagem direta em cada lote e concatena os logits em um único tensor na GPU. Assim que o conjunto de dados de avaliação é totalmente processado, o Trainer move os logits concatenados para a CPU e chama compute_metrics. Para modelos grandes ou conjuntos de avaliação extensos, os logits acumulados podem esgotar a memória da GPU, mesmo quando o treinamento no mesmo hardware funciona normalmente, pois o treinamento retém apenas um lote de ativações por vez.

eval_accumulation_steps

Desloque as previsões acumuladas da GPU para a CPU a cada n lotes. Valores menores reduzem a memória da GPU, mas aumentam a frequência das transferências para a CPU.

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

Chamado uma vez por lote de avaliação na GPU, imediatamente após a passagem para frente e antes da acumulação de logits. O valor retornado substitui os logits em eval_pred.predictions. Executar o cálculo no nível do lote reduz o tamanho do tensor por lote e fornece a eval_accumulation_steps um tensor menor para descarregar.

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

Desempenho do Dataloader

Por padrão, o Trainer cria um dataloader com dataloader_num_workers=0. Os dados são carregados no processo principal enquanto a GPU fica ociosa, o que se manifesta como baixa utilização da GPU entre os lotes.

Tanto dataloader_persistent_workers quanto dataloader_prefetch_factor exigem dataloader_num_workers > 0.

• dataloader_persistent_workers mantém os subprocessos de workers ativos entre épocas para evitar a reinitialização do zero, ao custo de maior uso de memória.
• dataloader_prefetch_factor controla quantos lotes cada worker prepara antecipadamente. Com dataloader_prefetch_factor=2 e num_workers=4, até 8 lotes ficam na memória enquanto a GPU treina com o lote atual.

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 amostras por comprimento

Use train_sampling_strategy="group_by_length" para agrupar exemplos com comprimentos semelhantes e reduzir o padding. Quando você não fornece comprimentos pré-computados, o Trainer os infere a partir da primeira entrada do modelo em cada item do dataset. Isso também funciona quando datasets multimodais baseados em processadores retornam objetos BatchFeature, pois eles são contêineres de recursos semelhantes a mapeamentos.

from transformers import TrainingArguments

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

Se um Dataset já possui uma coluna de comprimento pré-computada, o Trainer usa essa coluna. O nome padrão da coluna é length. Defina length_column_name quando o seu conjunto de dados usar outro nome. Essa estratégia requer um conjunto de dados com comprimento conhecido e é ignorado para IterableDataset.

Amostragem de reequilíbrio por lote

Em conjuntos de dados de comprimento variável, o desequilíbrio dentro de um micro-lote e entre dispositivos faz com que os dispositivos desperdicem tempo com preenchimento e fiquem ociosos nas etapas de sincronização de gradientes.

Defina train_sampling_strategy="batch_rebalance" em TrainingArguments para reduzir ambos os efeitos. Para cada passo do otimizador, o amostrador:

1. Ordena as amostras do lote por comprimento.
2. Distribui o lote ordenado entre os dispositivos de modo que o custo dos tokens de preenchimento de cada micro-lote seja equilibrado. Micro-lotes com amostras longas recebem menos amostras, e micro-lotes com amostras curtas recebem mais.

Isso reduz o padding dentro de cada micro-lote, e cada dispositivo finaliza um micro-lote em aproximadamente o mesmo tempo, o que reduz o tempo ocioso na sincronização e mantém a memória de pico menor do que com "group_by_length". Essa estratégia é suportada apenas para treinamento paralelo de dados por enquanto (o paralelismo de tensor ainda não é suportado).

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 é uma média aqui, em vez de uma contagem exata por etapa: alguns micro-lotes têm menos amostras e outros têm mais, mas o número total de amostras treinadas por etapa permanece o mesmo que com uma estratégia de amostragem distribuída normal.

O sampler precisa do comprimento de cada amostra para ordenar e equilibrar os lotes. Por padrão, o Trainer varre o conjunto de dados completo uma vez no início do treinamento para calcular esses comprimentos. Para pular essa varredura, pré-compute os comprimentos em uma coluna do conjunto de dados (por exemplo, durante o pré-processamento) e passe seu nome como length_column_name (padrão: "length").

NEFTune

NEFTune adiciona ruído aleatório aos embeddings de tokens durante a passagem direta. O ruído atua como regularização e pode melhorar o desempenho no ajuste fino por instrução.

Ative o NEFTune definindo neftune_noise_alpha em TrainingArguments. Os valores típicos de alpha variam 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()

O NEFTune afeta apenas o treinamento, e a camada de incorporação original é restaurada após o treinamento.

Registro

Controle quando e onde o Trainer grava as entradas de log com logging_strategy, logging_steps e report_to.

• logging_strategy="steps" registra a cada logging_steps() atualizações do otimizador. Use "epoch" para registrar no final de cada epoch, em vez disso.
• report_to envia logs para um 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()

Checkpointing

O Trainer salva um checkpoint a cada save_steps() atualização do otimizador e mantém todos eles (ou os mais recentes ~TrainingArguments.save_total_limit).

save_strategy="best" mantém apenas o único melhor checkpoint de acordo com uma métrica. Um novo checkpoint é salvo somente quando a métrica monitorada melhora, o que economiza espaço em disco e evita o acúmulo de checkpoints 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()

Retomar o treinamento

Passar resume_from_checkpoint=True para train() se o treinamento foi interrompido e você deseja retomar sem perder o progresso. O treinamento será retomado a partir do checkpoint mais recente em output_dir.

trainer.train(resume_from_checkpoint=True)

Especifique um caminho de checkpoint para retomar de um ponto específico.

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

Ao retomar, o Trainer restaura o estado do otimizador, o estado do agendador e o estado do RNG.

O retomar de checkpoints requer os arquivos de estado do otimizador e do agendador no diretório do checkpoint. Se esses arquivos estiverem ausentes (por exemplo, quando save_only_model=True), o otimizador será reiniciado do zero.

Checkpointing JIT

Com o checkpointing periódico (save_strategy="steps" ou "epoch"), você perde todo o progresso de treinamento entre o último checkpoint salvo e a interrupção. Em clusters compartilhados com cargas de trabalho preemptíveis, como o Kueue, os jobs podem ser encerrados a qualquer momento, de modo que esse intervalo pode significar horas de computação desperdiçada.

O checkpointing JIT (Just-In-Time) elimina essa lacuna. Quando o trainer recebe um sinal SIGTERM, ele salva um checkpoint no ponto exato em que o treinamento foi interrompido, permitindo que você retome com perda mínima de progresso. Ele funciona em conjunto com o checkpointing periódico. Os salvamentos periódicos protegem contra falhas e problemas de hardware, enquanto os salvamentos JIT protegem contra preempção e encerramentos graciosos.

Ative-o definindo enable_jit_checkpoint=True em TrainingArguments.

from transformers import TrainingArguments

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

Quando o SIGTERM é recebido, o Trainer aguarda a conclusão da etapa de treinamento atual, salva um checkpoint e interrompe o treinamento de forma graciosa. Um arquivo sentinela (checkpoint-is-incomplete.txt) é criado quando o salvamento começa e removido assim que o checkpoint é totalmente gravado. Se um diretório de checkpoint ainda contiver esse arquivo, o salvamento foi interrompido antes de ser concluído. O Trainer não verifica isso automaticamente, portanto, inspecione-o você mesmo antes de retomar.

Retome do checkpoint JIT da mesma forma que qualquer outro checkpoint.

trainer.train(resume_from_checkpoint=True)

Você deve configurar seu orquestrador para permitir tempo suficiente para que o checkpoint seja concluído. O período padrão de desligamento gracioso do Kubernetes é de apenas 30 segundos, o que geralmente não é suficiente para modelos maiores.

spec:
  template:
    spec:
      terminationGracePeriodSeconds: 300

#SBATCH --signal=TERM@300

Calcule o período de tolerância necessário como o tempo máximo possível de um passo de treinamento mais o tempo de salvamento do checkpoint, mais o atraso de 3 segundos de kill_wait antes do início do checkpoint. Por exemplo, se um passo de treinamento levar até 2 minutos e o salvamento de um checkpoint levar 2 minutos, defina pelo menos 243 segundos de tempo de tolerância.