huggingface-accelerate

accelerate launch train.py

# 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()

# 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()

accelerate config

# 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

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)

from accelerate import Accelerator
 
accelerator = Accelerator(
    mixed_precision='bf16',
    deepspeed_plugin={
        "zero_stage": 2,  # ZeRO-2
        "offload_optimizer": False,
        "gradient_accumulation_steps": 4
    }
)
 
# Same code as before!
model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)

accelerate config
# Select: DeepSpeed → ZeRO-2

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

accelerate launch --config_file deepspeed_config.json train.py

from accelerate import Accelerator, FullyShardedDataParallelPlugin
 
fsdp_plugin = FullyShardedDataParallelPlugin(
    sharding_strategy="FULL_SHARD",  # ZeRO-3 equivalent
    auto_wrap_policy="TRANSFORMER_AUTO_WRAP",
    cpu_offload=False
)
 
accelerator = Accelerator(
    mixed_precision='bf16',
    fsdp_plugin=fsdp_plugin
)
 
model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)

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

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()

# WRONG
batch = batch.to('cuda')
 
# CORRECT
# Accelerate handles it automatically after prepare()

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

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

from accelerate.utils import set_seed
set_seed(42)