Custom Training Loops with Accelerate

When Trainer is not enough

Trainer covers most supervised fine-tuning, but research ideas such as custom losses, unusual sampling or multiple models trained together are often clearer in a plain PyTorch loop. The pain of a custom loop is device handling: moving tensors to the right GPU, wrapping the model for distributed training and handling mixed precision. Accelerate removes that boilerplate. You create an Accelerator, pass your model, optimiser, dataloaders and scheduler through accelerator.prepare(...), and replace loss.backward() with accelerator.backward(loss). The same script then runs on a laptop CPU, a single GPU or several GPUs; for multi-GPU you configure once with accelerate config and start with accelerate launch train.py. Trainer itself is built on Accelerate, so both approaches share the same machinery. Use accelerator.print to print only once instead of once per process, and accelerator.gather_for_metrics to collect predictions from every process before computing metrics.

A minimal Accelerate loop

Assumes a tokenized dataset train_ds with a labels column and a data collator collator, for example DataCollatorWithPadding(tok) from the Trainer topic. No .to(device) calls are needed.

import torch
from torch.utils.data import DataLoader
from accelerate import Accelerator
from transformers import AutoModelForSequenceClassification, get_linear_schedule_with_warmup

accelerator = Accelerator(mixed_precision="bf16")  # or "no" / "fp16"
model = AutoModelForSequenceClassification.from_pretrained(
    "distilbert/distilbert-base-uncased", num_labels=2)
loader = DataLoader(train_ds, batch_size=16, shuffle=True, collate_fn=collator)
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)
scheduler = get_linear_schedule_with_warmup(optimizer, 100, len(loader) * 3)

model, optimizer, loader, scheduler = accelerator.prepare(model, optimizer, loader, scheduler)

model.train()
for epoch in range(3):
    for batch in loader:
        loss = model(**batch).loss
        accelerator.backward(loss)
        optimizer.step()
        scheduler.step()
        optimizer.zero_grad()
    accelerator.print(f"epoch {epoch} loss {loss.item():.4f}")

Remember zero_grad

PyTorch accumulates gradients by default. Forgetting optimizer.zero_grad() makes each step use the running sum of all previous gradients, so updates grow larger and larger and training usually diverges.

Quick check: In an Accelerate training loop, what replaces `loss.backward()`?

• optimizer.backward(loss)
• model.backward(loss)
• accelerator.backward(loss)
• torch.backward(loss, device="auto")