Distributed PyTorch training with HuggingFace Accelerate

HuggingFace Accelerate lets you run the same PyTorch training script on CPU, single or multi-GPU, multi-node, TPU, or Apple MPS setups by adding a few lines of code.

Hermes OpenClaw Intermediate Copied: 0
1 resource

Requirements

  • Requires the accelerate, torch, and transformers Python dependencies.
  • Accelerate is installed via pip.
  • DDP: torch.distributed.run (built-in)
  • DeepSpeed: deepspeed (pip install deepspeed)
  • FSDP is built into PyTorch 1.12 or later.
  • Megatron requires a custom setup.
  • Requires Accelerate version 1.11.0 or later.

How to use

Use Accelerate when you want the simplest distributed training, a single script for any hardware, the HuggingFace ecosystem, flexibility across DDP, DeepSpeed, FSDP, and Megatron, or quick prototyping.

It suits the simplest possible distributed training setup.

It supports a single script that runs on any hardware.

It fits well if you use the HuggingFace ecosystem.

It offers flexibility across DDP, DeepSpeed, FSDP, and Megatron.

It is suitable for quick prototyping.

Use alternatives instead when you need what they specialize in.

Choose PyTorch Lightning if you need callbacks and high-level abstractions.

Choose Ray Train for multi-node orchestration and hyperparameter tuning.

Choose DeepSpeed directly for API control and advanced features.

Choose raw DDP for maximum control with minimal abstraction.

CPU works, though slowly.

Single GPU works.

Multi-GPU runs with DDP (the default), DeepSpeed, or FSDP.

Multi-node supports DDP, DeepSpeed, FSDP, and Megatron.

TPU is supported.

Apple MPS is supported.

Install Accelerate to get started with the quick start.

Installation is done with pip.

pip install accelerate

Convert your PyTorch training script by adding four lines: create an Accelerator, call prepare on the model, optimizer, and dataloader, and replace loss.backward with accelerator.backward.

import torch
+ from accelerate import Accelerator

+ accelerator = Accelerator()

  model = torch.nn.Transformer()
  optimizer = torch.optim.Adam(model.parameters())
  dataloader = torch.utils.data.DataLoader(dataset)

+ model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)

  for batch in dataloader:
      optimizer.zero_grad()
      loss = model(batch)
-     loss.backward()
+     accelerator.backward(loss)
      optimizer.step()

Run the converted training script with a single launch command.

accelerate launch train.py

Start from an existing single-GPU PyTorch training script as your baseline.

# train.py
import torch

model = torch.nn.Linear(10, 2).to('cuda')
optimizer = torch.optim.Adam(model.parameters())
dataloader = torch.utils.data.DataLoader(dataset, batch_size=32)

for epoch in range(10):
    for batch in dataloader:
        batch = batch.to('cuda')
        optimizer.zero_grad()
        loss = model(batch).mean()
        loss.backward()
        optimizer.step()

With Accelerate, add four lines: import and instantiate the Accelerator, prepare the model, optimizer, and dataloader, and use accelerator.backward; manual .to('cuda') calls are no longer needed because device placement is automatic after prepare.

# train.py
import torch
from accelerate import Accelerator  # +1

accelerator = Accelerator()  # +2

model = torch.nn.Linear(10, 2)
optimizer = torch.optim.Adam(model.parameters())
dataloader = torch.utils.data.DataLoader(dataset, batch_size=32)

model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)  # +3

for epoch in range(10):
    for batch in dataloader:
        # No .to('cuda') needed - automatic!
        optimizer.zero_grad()
        loss = model(batch).mean()
        accelerator.backward(loss)  # +4
        optimizer.step()

Run the interactive configuration command to set up your environment.

accelerate config

Answer the interactive questions about your setup.

Choose which machine type you use: single or multi GPU, TPU, or CPU.

Specify how many machines, with 1 as the default.

Choose mixed precision: no, fp16, bf16, or fp8.

Choose whether to use DeepSpeed: no or yes.

Then launch the training; the same launch command works on any setup, including single GPU, multi-GPU with a process count, and multi-node with machine rank and main process IP.

# Single GPU
accelerate launch train.py

# Multi-GPU (8 GPUs)
accelerate launch --multi_gpu --num_processes 8 train.py

# Multi-node
accelerate launch --multi_gpu --num_processes 16 \
  --num_machines 2 --machine_rank 0 \
  --main_process_ip $MASTER_ADDR \
  train.py

Mixed precision training is enabled through the Accelerator constructor.

Set mixed_precision to fp16 (with gradient scaling), bf16 (no scaling, more stable), or fp8 (H100 GPUs and newer); the autocast context is optional and done automatically.

from accelerate import Accelerator

# FP16 (with gradient scaling)
accelerator = Accelerator(mixed_precision='fp16')

# BF16 (no scaling, more stable)
accelerator = Accelerator(mixed_precision='bf16')

# FP8 (H100+)
accelerator = Accelerator(mixed_precision='fp8')

model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)

# Everything else is automatic!
for batch in dataloader:
    with accelerator.autocast():  # Optional, done automatically
        loss = model(batch)
    accelerator.backward(loss)

Enable DeepSpeed ZeRO-2 by building a DeepSpeedPlugin and passing it to the Accelerator.

Pass a DeepSpeedPlugin instance, not a raw dict; the plugin sets the ZeRO stage, optimizer offload device, and gradient accumulation steps.

from accelerate import Accelerator, DeepSpeedPlugin

deepspeed_plugin = DeepSpeedPlugin(
    zero_stage=2,                     # ZeRO-2
    offload_optimizer_device="none",  # or "cpu" to offload
    gradient_accumulation_steps=4,
)

accelerator = Accelerator(
    mixed_precision='bf16',
    deepspeed_plugin=deepspeed_plugin,  # DeepSpeedPlugin instance (or dict[str, DeepSpeedPlugin])
)

# Same code as before!
model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)

Alternatively, point the plugin at a full DeepSpeed JSON config file, such as ds_config.json.

from accelerate import Accelerator, DeepSpeedPlugin

# hf_ds_config accepts a path to a DeepSpeed config JSON (or a dict)
deepspeed_plugin = DeepSpeedPlugin(hf_ds_config="ds_config.json")
accelerator = Accelerator(mixed_precision='bf16', deepspeed_plugin=deepspeed_plugin)

The ds_config.json file is a raw DeepSpeed config passed via the plugin, not via the launch command.

{
    "fp16": {"enabled": false},
    "bf16": {"enabled": true},
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {"device": "cpu"},
        "allgather_bucket_size": 5e8,
        "reduce_bucket_size": 5e8
    }
}

Alternatively, configure DeepSpeed through the interactive config by selecting DeepSpeed and ZeRO-2.

accelerate config
# Select: DeepSpeed → ZeRO-2
# This writes an accelerate YAML config (default: ~/.cache/huggingface/accelerate/default_config.yaml)

Launch training normally or point the config-file option at a specific accelerate YAML; the launch command expects an accelerate YAML config, not a raw DeepSpeed JSON.

# Uses the default accelerate config written by `accelerate config`
accelerate launch train.py

# Or point at a specific accelerate YAML
accelerate launch --config_file accelerate_deepspeed.yaml train.py

FSDP (Fully Sharded Data Parallel) is enabled through a FullyShardedDataParallelPlugin.

Build the plugin with a sharding strategy such as FULL_SHARD (the ZeRO-3 equivalent), an auto wrap policy, and a CPU offload setting, then pass it to the Accelerator.

from accelerate import Accelerator, FullyShardedDataParallelPlugin

fsdp_plugin = FullyShardedDataParallelPlugin(
    sharding_strategy="FULL_SHARD",  # ZeRO-3 equivalent
    auto_wrap_policy="transformer_based_wrap",  # valid: transformer_based_wrap | size_based_wrap | no_wrap
    cpu_offload=False
)

accelerator = Accelerator(
    mixed_precision='bf16',
    fsdp_plugin=fsdp_plugin
)

model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)

Alternatively, enable FSDP via the interactive config by selecting FSDP, Full Shard, and the CPU offload choice.

accelerate config
# Select: FSDP → Full Shard → No CPU Offload

Gradient accumulation is set on the Accelerator and applied with the accumulate context manager around your training loop.

Set gradient_accumulation_steps on the Accelerator, then wrap the loop body in the accumulate context manager, which handles the accumulation.

from accelerate import Accelerator

accelerator = Accelerator(gradient_accumulation_steps=4)

model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)

for batch in dataloader:
    with accelerator.accumulate(model):  # Handles accumulation
        optimizer.zero_grad()
        loss = model(batch)
        accelerator.backward(loss)
        optimizer.step()

Avoid manually moving batches to the device in distributed training.

Wrong device placement is a common issue.

Do not manually call batch.to('cuda'); Accelerate handles device placement automatically after prepare.

# WRONG
batch = batch.to('cuda')

# CORRECT
# Accelerate handles it automatically after prepare()

Gradient accumulation not working is a common issue.

Use the accelerator.accumulate context manager around the training steps so accumulation works.

# CORRECT
with accelerator.accumulate(model):
    optimizer.zero_grad()
    accelerator.backward(loss)
    optimizer.step()

Checkpointing in distributed training is a common issue area.

Use the accelerator methods: save state only on the main process and load state on all processes.

# Save only on main process
if accelerator.is_main_process:
    accelerator.save_state('checkpoint/')

# Load on all processes
accelerator.load_state('checkpoint/')

Getting different results with FSDP is a common issue.

Ensure the same random seed by setting it with set_seed.

from accelerate.utils import set_seed
set_seed(42)

DDP: torch.distributed.run (built-in)

DeepSpeed: deepspeed (pip install deepspeed)

Links

Source: nousresearch/hermes-agent and its contributors. License: MIT.

Copyright (c) 2025 Nous Research

Translation and editorial adaptation by ClawCases.

>Resources

>Related Use Cases

>More at this level: Intermediate