Fine-tuning with the Hugging Face Trainer

The Trainer — when you don't want to write the loop

The HuggingFace Trainer wraps the standard PyTorch training loop with batteries: distributed training, mixed precision, gradient accumulation, evaluation, checkpointing, logging, and a battle-tested set of defaults. For most fine-tuning tasks, it's faster to use Trainer than to roll your own loop.

The three pieces

The model — usually AutoModelFor[Task].from_pretrained(...).
The dataset — usually a HuggingFace Dataset, tokenized via .map(tokenizer, batched=True).
The training arguments — a TrainingArguments object encoding everything: epochs, batch size, LR, eval strategy, save strategy, FP16/BF16, logging dir.

What you give up vs raw PyTorch

The Trainer is opinionated. If your training loop has unusual structure (alternating two losses, custom gradient surgery, multi-stage curricula), you'll fight it. For "fine-tune model X on dataset Y", it's the right tool. For "train this novel architecture with this novel objective", roll your own loop.

When to use Trainer vs Lightning vs raw PyTorch

HF Trainer — best for fine-tuning HF models on HF datasets. The integration is tight.
PyTorch Lightning — best when you want loop boilerplate handled but the model isn't from HF, or you need more flexibility than Trainer gives.
Raw PyTorch — best for novel architectures, research, or when you want maximum control. The loop you wrote in Track 4 is enough for production.

Code

Fine-tune DistilBERT on IMDB — minimal example·python

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

model_name = "distilbert-base-uncased"
ds = load_dataset("imdb")
tok = AutoTokenizer.from_pretrained(model_name)

def tokenize(batch):
    return tok(batch["text"], padding="max_length", truncation=True, max_length=512)

ds_tok = ds.map(tokenize, batched=True)

model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)

args = TrainingArguments(
    output_dir="./out",
    num_train_epochs=3,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=64,
    learning_rate=2e-5,
    warmup_steps=500,
    weight_decay=0.01,
    eval_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True,
    bf16=True,                                 # bf16 mixed precision
    logging_dir="./logs",
    logging_steps=100,
)

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=ds_tok["train"],
    eval_dataset=ds_tok["test"],
    tokenizer=tok,                             # for proper save/load
)
trainer.train()

Custom metrics with compute_metrics·python

import numpy as np
from transformers import Trainer
from sklearn.metrics import accuracy_score, precision_recall_fscore_support

def compute_metrics(pred):
    labels = pred.label_ids
    preds = pred.predictions.argmax(-1)
    p, r, f1, _ = precision_recall_fscore_support(labels, preds, average='macro')
    return {
        'accuracy': accuracy_score(labels, preds),
        'precision': p,
        'recall': r,
        'f1': f1,
    }

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=ds_tok['train'],
    eval_dataset=ds_tok['test'],
    tokenizer=tok,
    compute_metrics=compute_metrics,
)

# Now eval prints these metrics each epoch
trainer.evaluate()

Resume training from a checkpoint·python

from transformers import Trainer

# Trainer automatically saves checkpoints to output_dir/checkpoint-N/
# Resume from a specific one:
trainer.train(resume_from_checkpoint="./out/checkpoint-1500")

# Or resume from the latest:
trainer.train(resume_from_checkpoint=True)

# Push the trained model to the HF Hub (requires `huggingface-cli login`)
# trainer.push_to_hub("my-fine-tuned-distilbert")

Exercise

Pick a small NLP dataset (rotten_tomatoes is ~10K samples, fast). Fine-tune distilbert-base-uncased for 1 epoch using Trainer. Add compute_metrics for accuracy. Verify the eval numbers print at end of training. Save the final model and load it with AutoModelForSequenceClassification.from_pretrained('./out').