HuggingFace Accelerate API idiom

Contract

Enforces the current HuggingFace Accelerate public API when you adapt a PyTorch training or evaluation script for distributed, multi-GPU, TPU, or mixed-precision execution. Apply whenever the task is to make PyTorch training run on more than one device, add mixed precision or gradient accumulation, do distributed evaluation, or write distributed-safe checkpoints with Accelerate.

Rules

1. Import and instantiate first. At the top of the script:

from accelerate import Accelerator, then accelerator = Accelerator(). One instance drives the whole script; create it before you prepare anything.

1. One prepare() wraps everything. Pass every training object — model, optimizer, each

dataloader (train and validation), and the LR scheduler — into a single accelerator.prepare(...) call, and unpack the return in the SAME order: model, optimizer, dataloader, scheduler = accelerator.prepare(model, optimizer, dataloader, scheduler). Do not reach for per-object helpers like prepare_model(...) / prepare_data_loader(...) in normal use — the single prepare() is the idiom.

1. Replace the backward pass. After prepare(), never call loss.backward(). Use

accelerator.backward(loss) — it scales gradients (for fp16) and picks the correct backend backward (DDP/DeepSpeed/Megatron).

1. Let Accelerate place tensors. Remove manual .to("cuda"), .cuda(), and .to(device) on the

model and on batches — prepare() and the prepared dataloader place them for you. If you genuinely need a device handle, use accelerator.device, never the literal string "cuda" or torch.device("cuda").

1. Mixed precision is a constructor kwarg, not a boolean.

Accelerator(mixed_precision="fp16") (or "bf16", or "fp8"). NOT Accelerator(fp16=True) (that argument was removed), and NOT a hand-built torch.cuda.amp.GradScaler + autocast. To cast a block computed OUTSIDE the model, wrap it in with accelerator.autocast():.

1. Gradient accumulation is a kwarg plus a context manager.

Accelerator(gradient_accumulation_steps=N), then wrap the per-step body in with accelerator.accumulate(model):. Do not hand-roll if (step + 1) % N == 0 with a manually divided loss.

1. Distributed evaluation gathers with gather_for_metrics. Before computing a metric across

processes, collect with accelerator.gather_for_metrics((predictions, targets)) — not raw torch.distributed.all_gather / dist.gather (which also leaves padding duplicates in).

1. Distributed-safe checkpointing.

• Full resumable state (model + optimizer + scheduler + RNG): accelerator.save_state(dir) and

accelerator.load_state(dir).

• Model-only export: call accelerator.wait_for_everyone(), then

unwrapped = accelerator.unwrap_model(model) and save unwrapped, guarding the write with if accelerator.is_main_process:. Never torch.save(model.state_dict()) on every rank without unwrapping first.

1. Gradient clipping goes through the accelerator. After prepare(), clip with

accelerator.clip_grad_norm_(...) (or clip_grad_value_), not torch.nn.utils.clip_grad_norm_.

1. Reproducibility across processes. from accelerate.utils import set_seed, then set_seed(42)

— it seeds every process/backend consistently. torch.manual_seed alone does not.

1. Print once. Use accelerator.print(...) (or guard with accelerator.is_main_process) so a

line is not duplicated once per process.

1. Launch with the CLI, not python. Run accelerate config once to generate the reusable

config (a default_config.yaml), then accelerate launch train.py. When not using a saved config, pass flags: --multi_gpu, --num_processes=N, --num_machines=N, --machine_rank=R, --main_process_ip=IP, --mixed_precision=fp16, --config_file=PATH. Do NOT document python train.py or python -m torch.distributed.launch as the launch path.

Worked examples

BEFORE = the base model's wrong default → AFTER = the conforming Accelerate idiom.

Core migration.

python
# BEFORE
device = "cuda"
model.to(device)
for batch in dataloader:
    batch = batch.to(device)
    optimizer.zero_grad()
    loss = model(batch).mean()
    loss.backward()
    optimizer.step()

# AFTER
from accelerate import Accelerator
accelerator = Accelerator()
model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)
for batch in dataloader:
    optimizer.zero_grad()
    loss = model(batch).mean()
    accelerator.backward(loss)
    optimizer.step()

Mixed precision.

python
# BEFORE
accelerator = Accelerator(fp16=True)          # removed argument
# BEFORE (hand-rolled)
scaler = torch.cuda.amp.GradScaler()
# AFTER
accelerator = Accelerator(mixed_precision="fp16")   # or "bf16" / "fp8"

Gradient accumulation.

python
# BEFORE
for i, batch in enumerate(dataloader):
    loss = model(batch).mean() / 4
    loss.backward()
    if (i + 1) % 4 == 0:
        optimizer.step(); optimizer.zero_grad()
# AFTER
accelerator = Accelerator(gradient_accumulation_steps=4)
for batch in dataloader:
    with accelerator.accumulate(model):
        loss = model(batch).mean()
        accelerator.backward(loss)
        optimizer.step(); optimizer.zero_grad()

Distributed metric gather.

python
# BEFORE
torch.distributed.all_gather(gathered, preds)
# AFTER
all_preds, all_targets = accelerator.gather_for_metrics((preds, targets))

Checkpointing.

python
# BEFORE
torch.save(model.state_dict(), "model.pt")   # every rank, wrapped model
# AFTER (resumable state)
accelerator.save_state("ckpt/")              # ... accelerator.load_state("ckpt/")
# AFTER (final export)
accelerator.wait_for_everyone()
if accelerator.is_main_process:
    torch.save(accelerator.unwrap_model(model).state_dict(), "model.pt")

Launch.

bash
# BEFORE
torchrun --nproc_per_node=8 train.py
# AFTER
accelerate config           # once, interactive
accelerate launch --multi_gpu --num_processes=8 train.py

Edge cases & exceptions

• Validation dataloader is prepared too. Pass the eval dataloader through prepare() as well, so

each process gets its shard; then use gather_for_metrics when scoring.

• save_state is called on ALL processes, not only the main one — it coordinates internally. Only

the model-export write and human-facing logging are guarded by is_main_process.

• bf16 needs no scaler; fp16 scaling is automatic through accelerator.backward — you never build

a GradScaler yourself.

• autocast() is usually unnecessary because backward already handles the loss; use

with accelerator.autocast(): only for extra mixed-precision work outside the model.

• The backend is unified. Switching to DeepSpeed or FSDP does NOT change the loop — prepare()

and accelerator.backward(loss) stay identical; you select the backend via accelerate config (or a plugin passed to Accelerator), not by rewriting the training code.

• torchrun still works, but the documented Accelerate path is accelerate launch; prefer it.

Do / Don't

• DON'T call loss.backward() after prepare(). ALWAYS call accelerator.backward(loss).
• DON'T leave .to("cuda") / .cuda() on model or batches. ALWAYS let prepare() place them (use

accelerator.device if you truly need one).

• DON'T write Accelerator(fp16=True). ALWAYS write Accelerator(mixed_precision="fp16").
• DON'T hand-roll accumulation with % N. ALWAYS use gradient_accumulation_steps +

with accelerator.accumulate(model):.

• DON'T gather metrics with torch.distributed.all_gather. ALWAYS use gather_for_metrics.
• DON'T torch.save(model.state_dict()) on every rank. ALWAYS unwrap_model + is_main_process, or

save_state/load_state.

• DON'T launch via python -m torch.distributed.launch. ALWAYS use accelerate launch (after

accelerate config).

Common mistakes (the base model's wrong defaults)

• Keeping loss.backward() after wrapping the model.
• Leaving model.to("cuda") and batch.to("cuda") in the loop.
• Enabling fp16 with the removed fp16=True flag or a manual GradScaler.
• Emulating gradient accumulation with a step-modulo check and a hand-divided loss.
• Using torch.distributed.all_gather for validation metrics.
• Saving the wrapped model with torch.save on all ranks.
• Launching with python train.py or python -m torch.distributed.launch.
• Seeding with only torch.manual_seed.

Quick checklist

• ] from accelerate import Accelerator + accelerator = Accelerator() at the top.
• ] One accelerator.prepare(...) for model, optimizer, all dataloaders, scheduler.
• ] accelerator.backward(loss) instead of loss.backward().
• ] No manual .to("cuda")/.cuda(); use accelerator.device if needed.
• ] mixed_precision="fp16"|"bf16"|"fp8" kwarg (not fp16=True).
• ] gradient_accumulation_steps + with accelerator.accumulate(model):.
• ] gather_for_metrics for distributed eval.
• ] save_state/load_state, or unwrap_model + is_main_process for export.
• ] set_seed from accelerate.utils; accelerator.print for logs.
• ] accelerate config then accelerate launch (with --multi_gpu/--num_processes/…).