Распределенное обучение в PyTorch с помощью HuggingFace Accelerate

Как применить

Используйте Accelerate, если нужно самое простое распределенное обучение, один скрипт для любого оборудования, экосистема HuggingFace, гибкость при использовании DDP, DeepSpeed, FSDP и Megatron, или быстрое прототипирование.

Подходит для самой простой возможной конфигурации распределенного обучения.

Поддерживает один скрипт, который работает на любом оборудовании.

Хорошо подходит, если вы используете экосистему HuggingFace.

Обеспечивает гибкость при работе с DDP, DeepSpeed, FSDP и Megatron.

Подходит для быстрого прототипирования.

Используйте альтернативные решения, если вам нужны их специализированные возможности.

Выберите PyTorch Lightning, если вам нужны коллбэки и высокоуровневые абстракции.

Выберите Ray Train для оркестрации в многоузловых средах и настройки гиперпараметров.

Выберите DeepSpeed напрямую для управления API и расширенных функций.

Выберите чистый DDP для максимального контроля при минимальных абстракциях.

Работает на CPU, хотя и медленно.

Работает на одной GPU.

Много-GPU работает с DDP (по умолчанию), DeepSpeed или FSDP.

Многоузловая работа поддерживает DDP, DeepSpeed, FSDP и Megatron.

Поддерживается TPU.

Поддерживается Apple MPS.

Установите Accelerate, чтобы начать работу с быстрого старта.

Установка осуществляется через pip.

pip install accelerate

Преобразуйте свой скрипт обучения PyTorch, добавив четыре строки: создайте объект Accelerator, вызовите метод prepare для модели, оптимизатора и загрузчика данных, а также замените loss.backward на 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()

Запустите преобразованный скрипт обучения одной командой запуска.

accelerate launch train.py

Начните с существующего скрипта обучения PyTorch на одной GPU в качестве базы.

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

С помощью Accelerate добавьте четыре строки: импортируйте и создавайте экземпляр Accelerator, подготовьте модель, оптимизатор и загрузчик данных, затем используйте accelerator.backward; явные вызовы .to('cuda') больше не нужны, так как размещение устройства происходит автоматически после вызова 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()

Запустите интерактивную команду конфигурации для настройки вашей среды.

accelerate config

Ответьте на интерактивные вопросы о вашей конфигурации.

Выберите тип используемой машины: одна или несколько GPU, TPU или CPU.

Укажите количество машин; по умолчанию — 1.

Выберите режим смешанной точности: нет, fp16, bf16 или fp8.

Выберите, использовать ли DeepSpeed: нет или да.

Затем запустите обучение; одна и та же команда запуска работает на любой конфигурации, включая одну GPU, много-GPU с указанием числа процессов и многоузловую среду с указанием номера узла и 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

Режим смешанной точности включается через конструктор Accelerator.

Установите параметр mixed_precision в значение fp16 (с масштабированием градиентов), bf16 (без масштабирования, более стабильно) или fp8 (для GPU H100 и новее); контекст autocast является необязательным и применяется автоматически.

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)

Включите DeepSpeed ZeRO-2, создав объект DeepSpeedPlugin и передав его в Accelerator.

Передавайте экземпляр DeepSpeedPlugin, а не сырой словарь; плагин задает этап ZeRO, устройство для выгрузки оптимизатора и количество шагов накопления градиентов.

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)

В качестве альтернативы укажите плагину полный файл конфигурации DeepSpeed в формате JSON, например 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)

Файл ds_config.json представляет собой сырую конфигурацию DeepSpeed, передаваемую через плагин, а не через команду запуска.

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

Как вариант, настройте DeepSpeed через интерактивный конфигуратор, выбрав DeepSpeed и режим ZeRO-2.

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

Запустите обучение в обычном режиме или укажите параметр config-file на конкретный файл конфигурации accelerate в формате YAML; команда запуска ожидает именно такой файл, а не сырой JSON-конфиг DeepSpeed.

# 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) включается с помощью плагина FullyShardedDataParallelPlugin.

Соберите плагин с такой стратегией шардирования, как FULL_SHARD (эквивалент ZeRO-3), политикой автоматического обёртывания и настройкой выгрузки на CPU, затем передайте его в 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)

Как вариант, включите режим FSDP через интерактивный конфигуратор, выбрав FSDP, режим Full Shard и вариант выгрузки на CPU.

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

Накопление градиентов настраивается в Accelerator и применяется с помощью менеджера контекста accumulate вокруг цикла обучения.

Установите параметр gradient_accumulation_steps в Accelerator, затем оберните тело цикла в менеджер контекста accumulate — он и будет отвечать за накопление.

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

Избегайте ручного перемещения батчей на устройство при распределённом обучении.

Неправильное размещение устройств — распространённая проблема.

Не вызывайте функцию batch.to('cuda') вручную; Accelerate автоматически управляет размещением устройств после вызова prepare.

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

# CORRECT
# Accelerate handles it automatically after prepare()

Неработающий механизм накопления градиентов — частая проблема.

Используйте менеджер контекста accelerator.accumulate вокруг шагов обучения, чтобы накопление градиентов функционировало корректно.

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

Проблемы с сохранением точек восстановления при распределённом обучении — распространённая область сложностей.

Используйте методы Accelerator: сохраняйте состояние только на главном процессе и загружайте его на всех остальных процессах.

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

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

Получение разных результатов при использовании FSDP — частая проблема.

Чтобы гарантировать одинаковые случайные числа, установите одинаковое значение семени с помощью функции set_seed.

from accelerate.utils import set_seed
set_seed(42)

DDP: torch.distributed.run (встроенный)

DeepSpeed: deepspeed (pip install deepspeed)

Ссылки

• references/megatron-integration.md
• references/custom-plugins.md
• references/performance.md