Hugging Face's logo

Accelerate documentation

DDP Communication Hooks

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v1.15.0).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

DDP Communication Hooks

Хуки коммуникации Distributed Data Parallel (DDP) предоставляют универсальный интерфейс для управления тем, как градиенты передаются между рабочими процессами, переопределяя стандартный allreduce в DistributedDataParallel. Предусмотрено несколько встроенных хуков коммуникации, и пользователи могут легко применить любой из них для оптимизации передачи данных.

  • Хук сжатия FP16: сжимает градиенты, приводя их к формату чисел с плавающей точкой половинной точности (torch.float16), что снижает накладные расходы на коммуникацию.
  • BF16 Compression Hook: Аналогично FP16, но использует формат Brain Floating Point (torch.bfloat16), который может быть более эффективным на определённом оборудовании.
  • PowerSGD Hook: продвинутый алгоритм сжатия градиентов, обеспечивающий высокие коэффициенты сжатия и способный ускорять распределённое обучение, ограниченное пропускной способностью.

В этом руководстве вы узнаете, как быстро настроить хуки коммуникации DDP и выполнить обучение с помощью утилит, предоставляемых в Accelerate, что может сводиться к добавлению всего одной новой строки кода! Это демонстрирует, как использовать хуки коммуникации DDP для оптимизации передачи градиентов в распределённом обучении с помощью библиотеки Accelerate.

FP16 Compression Hook

PyTorch
Accelerate
import torch
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed.algorithms.ddp_comm_hooks import default_hooks
from accelerate.test_utils.testing import get_backend

device_type, _, _ = get_backend()
device_id = getattr(torch, device_type, torch.cuda).current_device()

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(10, 10)

    def forward(self, x):
        return self.layer(x)

model = MyModel()
model = DDP(model, device_ids=[device_id])
model.register_comm_hook(state=None, hook=default_hooks.fp16_compress_hook)

# Training loop
for data, targets in data_loader:
    outputs = model(data)
    loss = criterion(outputs, targets)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

BF16 Compression Hook

API BF16 Compression Hook является экспериментальным и требует версию NCCL, более позднюю, чем 2.9.6.

PyTorch
Accelerate
import torch
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed.algorithms.ddp_comm_hooks import default_hooks
from accelerate.test_utils.testing import get_backend

device_type, _, _ = get_backend()
device_id = getattr(torch, device_type, torch.cuda).current_device()

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(10, 10)

    def forward(self, x):
        return self.layer(x)

model = MyModel()
model = DDP(model, device_ids=[device_id])
model.register_comm_hook(state=None, hook=default_hooks.bf16_compress_hook)

# Training loop
for data, targets in data_loader:
    outputs = model(data)
    loss = criterion(outputs, targets)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

PowerSGD Hook

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

PyTorch
Accelerate
import torch
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed.algorithms.ddp_comm_hooks import powerSGD_hook
from accelerate.test_utils.testing import get_backend

device_type, _, _ = get_backend()
device_id = getattr(torch, device_type, torch.cuda).current_device()

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(10, 10)

    def forward(self, x):
        return self.layer(x)

model = MyModel()
model = DDP(model, device_ids=[device_id])
state = powerSGD_hook.PowerSGDState(process_group=None)
model.register_comm_hook(state=state, hook=powerSGD_hook.powerSGD_hook)

# Training loop
for data, targets in data_loader:
    outputs = model(data)
    loss = criterion(outputs, targets)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

Утилиты для хуков коммуникации DDP

Существуют два дополнительных утилитарных средства для поддержки необязательных функций с хуками коммуникации.

comm_wrapper

comm_wrapper — это параметр для обёртки коммуникационного хука с дополнительной функциональностью. Например, его можно использовать для объединения сжатия FP16 с другими стратегиями коммуникации. В настоящее время поддерживаются обёртки no, fp16 и bf16.

from accelerate import Accelerator, DDPCommunicationHookType, DistributedDataParallelKwargs
import torch

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(10, 10)

    def forward(self, x):
        return self.layer(x)

# DDP Communication Hook setup
ddp_kwargs = DistributedDataParallelKwargs(
    comm_hook=DDPCommunicationHookType.POWER_SGD,
    comm_wrapper=DDPCommunicationHookType.FP16
)
accelerator = Accelerator(kwargs_handlers=[ddp_kwargs])

model = MyModel()
optimizer = torch.optim.Adam(model.parameters())
data_loader = DataLoader(dataset, batch_size=16)

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

# Training loop
for data, targets in data_loader:
    outputs = model(data)
    loss = criterion(outputs, targets)
    accelerator.backward(loss)
    optimizer.step()
    optimizer.zero_grad()

comm_state_option

comm_state_option позволяет передавать дополнительную информацию о состоянии, необходимую для определённых хуков коммуникации. Это особенно полезно для хуков с состоянием, таких как PowerSGD, которые требуют сохранения гиперпараметров и внутренних состояний между шагами обучения. Ниже приведён пример использования comm_state_option с хуком PowerSGD.

from accelerate import Accelerator, DDPCommunicationHookType, DistributedDataParallelKwargs
import torch

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(10, 10)

    def forward(self, x):
        return self.layer(x)

# DDP Communication Hook setup
ddp_kwargs = DistributedDataParallelKwargs(
    comm_hook=DDPCommunicationHookType.POWER_SGD,
    comm_state_option={"matrix_approximation_rank": 2}
)
accelerator = Accelerator(kwargs_handlers=[ddp_kwargs])

model = MyModel()
optimizer = torch.optim.Adam(model.parameters())
data_loader = DataLoader(dataset, batch_size=16)

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

# Training loop
for data, targets in data_loader:
    outputs = model(data)
    loss = criterion(outputs, targets)
    accelerator.backward(loss)
    optimizer.step()
    optimizer.zero_grad()

Для более продвинутого использования и дополнительных хуков обратитесь к документации PyTorch DDP Communication Hooks.

Update on GitHub

←Использование нескольких моделей с DeepSpeed Полностью шардированная параллельная обработка данных→