C.W.K.
Stream
Lesson 01 of 08 · published

Trainer and TrainingArguments

~30 min · training, trainer

Level 0Scout
0 XP0/50 lessons0/10 achievements
0/120 XP to next level120 XP to go0% complete

The Trainer pattern

transformers.Trainer is the canonical training loop. You assemble: a model, a tokenizer, a tokenized Dataset, optional metrics, optional callbacks, and a TrainingArguments object. trainer.train() runs; trainer.evaluate() reports; trainer.save_model() persists. The same loop handles fine-tuning, LoRA training, and full pre-training (with appropriate args).

TrainingArguments knobs you'll set 95% of the time

  • output_dir — checkpoints + tensorboard logs.
  • num_train_epochs or max_steps — budget.
  • per_device_train_batch_size + gradient_accumulation_steps — effective batch.
  • learning_rate, weight_decay, warmup_ratio, lr_scheduler_type — optimization.
  • bf16 / fp16 — mixed precision.
  • eval_strategy, eval_steps, save_strategy, save_steps — cadence.
  • logging_steps, report_to=['tensorboard', 'wandb'] — observability.

What Trainer hides

Distributed setup (single-node multi-GPU "just works" via accelerate), gradient accumulation, mixed precision, gradient clipping, learning-rate scheduling, checkpoint resumption. The price: it's opinionated. For exotic optimizers / non-standard losses, drop down to the raw PyTorch loop — or use trl / peft Trainer subclasses.

Code

Smallest possible Trainer-based fine-tune·python
from transformers import (
    AutoTokenizer, AutoModelForSequenceClassification,
    Trainer, TrainingArguments, DataCollatorWithPadding,
)
from datasets import load_dataset
import numpy as np

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

ds = load_dataset("stanfordnlp/imdb")
def tokenize(b): return tok(b["text"], truncation=True, max_length=256)
ds_tok = ds.map(tokenize, batched=True, remove_columns=["text"])

args = TrainingArguments(
    output_dir="./out",
    num_train_epochs=1,
    per_device_train_batch_size=16,
    learning_rate=2e-5,
    warmup_ratio=0.1,
    weight_decay=0.01,
    bf16=True,
    eval_strategy="steps",
    eval_steps=200,
    save_strategy="steps",
    save_steps=500,
    logging_steps=50,
    report_to=["tensorboard"],
)

def metrics(pred):
    preds = pred.predictions.argmax(-1)
    return {"accuracy": (preds == pred.label_ids).mean()}

trainer = Trainer(
    model=model, args=args,
    train_dataset=ds_tok["train"].select(range(2000)),
    eval_dataset=ds_tok["test"].select(range(500)),
    tokenizer=tok,
    data_collator=DataCollatorWithPadding(tok),
    compute_metrics=metrics,
)
trainer.train()
print(trainer.evaluate())

External links

Exercise

Run the smallest-possible Trainer fine-tune above on distilbert + imdb. Verify TensorBoard logs show loss decreasing. Check the saved checkpoint dir contains weights, optimizer state, scheduler state, and trainer_state.json. Resume training from the checkpoint and verify it picks up where it left off.

Progress

Progress is local-only — sign in to sync across devices.
Spotted a bug or have feedback on this page?Report an Issue

Comments 0

🔔 Reply notifications (sign in)
Sign inPlease sign in to comment.

No comments yet — be the first.