accelerate
Accelerate is Hugging Face's thin wrapper that lets a plain PyTorch training loop run on any hardware setup: single GPU, multi-GPU, multi-node, TPU or CPU, with fp16/bf16/fp8 mixed precision, DeepSpeed and FSDP switched on through configuration instead of code rewrites. You add roughly five lines (create an Accelerator, prepare() your model, optimizer and dataloaders, replace loss.backward() with accelerator.backward(loss)) and start jobs with the accelerate CLI. It is also the machinery behind device_map='auto' big-model loading in transformers.
The right abstraction if you want to keep writing your own training loop: one class, config-driven scaling, and structural Hugging Face support since transformers Trainer runs on it. Respect the porting checklist, because the failure mode is code that runs cleanly while computing the wrong thing.
Use it if
- You write your own PyTorch training loops and want them to scale from a laptop to multi-GPU or multi-node without learning torch.distributed launch incantations
- You want to switch between DDP, DeepSpeed ZeRO and FSDP by editing a config file rather than restructuring training code
- You need big-model inference: loading checkpoints larger than one GPU with automatic device maps and CPU or disk offload is an Accelerate feature
- You train from notebooks: notebook_launcher runs distributed training from Colab or Kaggle cells where a CLI launcher is unavailable
- You do not want to own a training loop at all: the README itself says so, and PyTorch Lightning or the transformers Trainer give you loops, callbacks and logging out of the box
- You will only ever train on one GPU: plain PyTorch works fine and is one less abstraction between you and a debugger
- You expect it to make distributed debugging easy: it removes launch boilerplate, not NCCL timeouts, hanging collectives or DeepSpeed config sharp edges, which still require understanding the stack underneath
- Your team will not learn the porting rules: gradient accumulation contexts, gather_for_metrics, main-process guards; loops ported carelessly run without errors while training or evaluating subtly wrong
Setup reality
pip install accelerate is light, but PyTorch must already be installed for your CUDA version, and anything beyond DDP means extra installs; deepspeed in particular compiles CUDA extensions and fails in creative ways. The workflow is accelerate config once per machine (it writes a YAML answers file) then accelerate launch script.py, and forgetting that a stale config exists on a box is a classic source of mystery behavior. Porting rules that bite: every dataloader must pass through prepare() or data is not sharded across processes, checkpoints should go through save_state or unwrap_model before save_pretrained, metrics need gather_for_metrics to drop duplicated samples from the padded last batch, and prints belong behind is_main_process unless you want one line per GPU.
Patterns
Convert a PyTorch loop to Accelerateminimal-training-loop
from accelerate import Accelerator
accelerator = Accelerator()
model, optimizer, loader = accelerator.prepare(model, optimizer, loader)
model.train()
for epoch in range(10):
for batch in loader:
optimizer.zero_grad()
loss = compute_loss(model, batch)
accelerator.backward(loss)
optimizer.step()Do not call .to(device) or .cuda() yourself after prepare(); device placement is handled, and batches from a prepared dataloader already sit on the right device.
Configure once, launch anywhereconfigure-and-launch
accelerate config # interactive, writes default_config.yaml
accelerate launch train.py --my-args
# skip the config file entirely
accelerate launch --multi_gpu --num_processes 2 train.py
# sanity-check the environment
accelerate envaccelerate launch reads the saved config by default, so behavior differs per machine until you pass explicit flags; plain python train.py still works single-process.
Train in bf16 or fp16mixed-precision
accelerator = Accelerator(mixed_precision='bf16') # or 'fp16', 'fp8'
model, optimizer, loader = accelerator.prepare(model, optimizer, loader)
for batch in loader:
with accelerator.autocast():
loss = compute_loss(model, batch)
accelerator.backward(loss)
optimizer.step()
optimizer.zero_grad()accelerator.backward handles fp16 loss scaling for you; prefer bf16 on Ampere or newer GPUs since it needs no scaler and rarely diverges. fp8 additionally requires TransformerEngine or MS-AMP.
Gradient accumulation without manual bookkeepinggradient-accumulation
accelerator = Accelerator(gradient_accumulation_steps=4)
model, optimizer, loader = accelerator.prepare(model, optimizer, loader)
for batch in loader:
with accelerator.accumulate(model):
loss = compute_loss(model, batch)
accelerator.backward(loss)
optimizer.step()
optimizer.zero_grad()Inside accumulate(), optimizer.step() only really steps every N batches and gradient sync is skipped on the others, saving communication; do not also divide the loss by N.
Checkpoint and resume full training statesave-and-resume-checkpoint
# save model, optimizer, scheduler, RNG and scaler state
accelerator.save_state('ckpt/step_1000')
# later, after building and preparing the same objects
accelerator.load_state('ckpt/step_1000')
# export just the model for inference
accelerator.wait_for_everyone()
unwrapped = accelerator.unwrap_model(model)
unwrapped.save_pretrained('out', save_function=accelerator.save,
state_dict=accelerator.get_state_dict(model))save_state is for resuming; unwrap_model strips the DDP/DeepSpeed wrapper for export, and get_state_dict gathers sharded ZeRO-3 weights onto one process.
Evaluate correctly across processesdistributed-evaluation
model.eval()
all_preds, all_labels = [], []
for batch in eval_loader: # must also be prepare()d
with torch.no_grad():
logits = model(**batch).logits
preds, labels = accelerator.gather_for_metrics(
(logits.argmax(dim=-1), batch['labels'])
)
all_preds.append(preds)
all_labels.append(labels)gather_for_metrics drops the duplicated samples that pad the last batch across processes; plain gather() does not, which quietly inflates your metrics.
Guard logging and side effectsmain-process-logging
accelerator.print('only prints once') # instead of print
if accelerator.is_main_process:
save_metrics(metrics)
accelerator.wait_for_everyone() # barrier before shared-file accessEvery process runs your whole script; unguarded writes race and unguarded logs multiply by the process count. is_local_main_process exists for per-node work like downloads.
Load a model too big for one GPUbig-model-inference
from accelerate import init_empty_weights, load_checkpoint_and_dispatch
with init_empty_weights():
model = MyModel(config) # no memory allocated
model = load_checkpoint_and_dispatch(
model, 'checkpoint_dir', device_map='auto',
offload_folder='offload',
)device_map='auto' spreads layers across GPUs, then CPU RAM, then disk; this is inference-only sharding, not training parallelism, so throughput drops once offload kicks in.
Launch distributed training from a notebooknotebook-launcher
from accelerate import notebook_launcher
def training_function():
accelerator = Accelerator()
...
notebook_launcher(training_function, num_processes=2)Nothing in the notebook may touch CUDA before the launcher spawns workers, or you get a 'CUDA has been initialized' error; restart the kernel and keep setup inside the function.
Enable DeepSpeed ZeRO from codedeepspeed-zero
from accelerate import Accelerator, DeepSpeedPlugin
plugin = DeepSpeedPlugin(zero_stage=2, gradient_accumulation_steps=2)
accelerator = Accelerator(mixed_precision='fp16', deepspeed_plugin=plugin)
model, optimizer, loader = accelerator.prepare(model, optimizer, loader)Most people set this through accelerate config instead; the plugin must know gradient_accumulation_steps up front, and ZeRO-3 changes how you must save weights (get_state_dict).
Seed every process consistentlyreproducible-seeding
from accelerate.utils import set_seed
set_seed(42)
# per-process different seeds for data augmentation:
set_seed(42, device_specific=True)set_seed covers python, numpy and torch on all processes; device_specific offsets by process index so augmentations differ per worker while staying reproducible.
Alternatives
| Package | Registry | Pick it when |
|---|---|---|
| pytorch-lightning | PyPI | A full training framework with loops, callbacks and loggers when you would rather not write the loop yourself |
| deepspeed | PyPI | Direct DeepSpeed usage when you need fine control over ZeRO stages and offload beyond what the plugin exposes |
| ray | PyPI | Ray Train when scaling across a cluster with scheduling, fault tolerance and hyperparameter search attached |