Fine-tuning with Trainer

Trainer gives you a correct, feature-rich training loop (mixed precision, multi-GPU, logging) without writing one.

Trainer implements the full PyTorch training loop for Transformers models: batching, forward and backward passes, optimizer and LR schedule (AdamW with linear decay by default), evaluation, checkpointing, logging and multi-device support.

You configure it with TrainingArguments:

output_dir: where checkpoints go.
learning_rate (2e-5 to 5e-5 is typical for full fine-tuning), num_train_epochs, per_device_train_batch_size, weight_decay, warmup_steps.
eval_strategy="epoch" or "steps" and a matching save_strategy; load_best_model_at_end=True with metric_for_best_model.
logging_steps, report_to="none" or "tensorboard", "wandb", "mlflow".
bf16=True on recent GPUs for mixed precision.

Then:

trainer = Trainer(model=model, args=args, train_dataset=train, eval_dataset=val,
                  processing_class=tok, data_collator=collator)
trainer.train()

processing_class (the tokenizer or processor) is saved alongside checkpoints so they load as complete models. trainer.evaluate() returns a dict of metrics such as eval_loss; trainer.predict(ds) returns predictions, label ids and metrics.

Resume an interrupted run with trainer.train(resume_from_checkpoint=True). The model must return a loss when given labels, which all AutoModelFor... heads do.

Example

from datasets import load_dataset
from transformers import (AutoModelForSequenceClassification, AutoTokenizer,
                          DataCollatorWithPadding, Trainer, TrainingArguments)

name = "distilbert/distilbert-base-uncased"
tok = AutoTokenizer.from_pretrained(name)
model = AutoModelForSequenceClassification.from_pretrained(name, num_labels=2)

ds = load_dataset("stanfordnlp/imdb")
ds = ds.map(lambda b: tok(b["text"], truncation=True, max_length=256), batched=True)
train = ds["train"].shuffle(seed=42).select(range(2000))
val = ds["test"].shuffle(seed=42).select(range(500))

args = TrainingArguments(
    output_dir="imdb-distilbert", learning_rate=2e-5, num_train_epochs=2,
    per_device_train_batch_size=16, eval_strategy="epoch", save_strategy="epoch",
    load_best_model_at_end=True, logging_steps=50, report_to="none",
)
trainer = Trainer(model=model, args=args, train_dataset=train, eval_dataset=val,
                  processing_class=tok, data_collator=DataCollatorWithPadding(tok))
trainer.train()
print(trainer.evaluate())   # {'eval_loss': 0.3..., 'epoch': 2.0, ...}

Remember

TrainingArguments configures; Trainer runs the loop
Pass processing_class=tokenizer so checkpoints include it
Match eval_strategy and save_strategy when using load_best_model_at_end

Watch out for

Using the removed argument name evaluation_strategy; current versions only accept eval_strategy.

Check yourself

You set load_best_model_at_end=True with eval_strategy='epoch' and save_strategy='steps'. What happens?

Official docs: huggingface.co/docs/transformers/training

Example

Remember

Watch out for

Check yourself

Builds on

Unlocks