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

Các hook giao tiếp Distributed Data Parallel (DDP) cung cấp một giao diện chung để kiểm soát cách các gradient được truyền thông giữa các worker bằng cách ghi đè allreduce mặc định trong DistributedDataParallel. Một số hook giao tiếp tích hợp sẵn được cung cấp, và người dùng có thể dễ dàng áp dụng bất kỳ hook nào trong số này để tối ưu hóa giao tiếp.

  • FP16 Compression Hook: Nén gradient bằng cách chuyển chúng sang định dạng số dấu phẩy động độ chính xác nửa (torch.float16), giảm thiểu chi phí truyền thông.
  • BF16 Compression Hook: Tương tự như FP16, nhưng sử dụng định dạng Brain Floating Point (torch.bfloat16), có thể hiệu quả hơn trên một số phần cứng nhất định.
  • PowerSGD Hook: Một thuật toán nén gradient nâng cao, cung cấp tỷ lệ nén cao và có thể tăng tốc độ huấn luyện phân phối bị giới hạn bởi băng thông.

Trong hướng dẫn này, bạn sẽ thấy cách thiết lập nhanh các hook giao tiếp DDP và thực hiện huấn luyện bằng các tiện ích có sẵn trong Accelerate, có thể đơn giản chỉ là thêm một dòng code mới! Điều này minh họa cách sử dụng các hook giao tiếp DDP để tối ưu hóa việc truyền thông gradient trong huấn luyện phân tán với thư viện 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()

Hook Nén BF16

API Hook Nén BF16 mang tính thử nghiệm và yêu cầu phiên bản NCCL mới hơn 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 thường yêu cầu bộ nhớ bổ sung có kích thước tương đương với gradient của mô hình để kích hoạt phản hồi lỗi, giúp bù đắp cho việc truyền thông bị nén có thiên lệch và cải thiện độ chính xác.

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

Công cụ Hooks liên lạc DDP

Có hai tiện ích bổ sung để hỗ trợ các tính năng tùy chọn với các hook truyền thông.

comm_wrapper

comm_wrapper là một tùy chọn để bọc một hook truyền thông với các chức năng bổ sung. Ví dụ, nó có thể được sử dụng để kết hợp nén FP16 với các chiến lược truyền thông khác. Các wrapper được hỗ trợ hiện tại là no, fp16 và 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 cho phép bạn truyền thêm thông tin trạng thái cần thiết cho một số hook truyền thông nhất định. Điều này đặc biệt hữu ích cho các hook có trạng thái như PowerSGD, vốn yêu cầu duy trì các siêu tham số và trạng thái nội bộ xuyên suốt các bước huấn luyện. Dưới đây là một ví dụ minh họa việc sử dụng comm_state_option cùng với hook 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()

Để biết thêm về cách sử dụng nâng cao và các hook bổ sung, hãy tham khảo tài liệu PyTorch DDP Communication Hooks.

Update on GitHub

←Sử dụng nhiều mô hình với DeepSpeed Song song dữ liệu phân mảnh hoàn toàn→