# Custom Training Loops with Accelerate — Hugging Face Transformers

Source: https://www.skillbyai.com/en/hugging-face/hf-accelerate-loop

> Write your own PyTorch training loop that runs unchanged on CPU, one GPU or several.

## 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.

```python
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.

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

- [ ] optimizer.backward(loss)
- [ ] model.backward(loss)
- [x] accelerator.backward(loss)
- [ ] torch.backward(loss, device="auto")

*Answer:* accelerator.backward(loss). accelerator.backward handles scaling for mixed precision and distributed gradient synchronisation.