Accelerate documentation

DDP 통신 훅

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 통신 훅

Distributed Data Parallel(DDP) 통신 훅은 DistributedDataParallel의 기본 allreduce를 오버라이드하여 워커 간 기울기 통신 방식을 제어하는 범용 인터페이스를 제공합니다. 몇 가지 내장된 통신 훅이 제공되며, 사용자는 이러한 훅 중 하나를 쉽게 적용하여 통신을 최적화할 수 있습니다.

  • FP16 압축 훅: 그래디언트를 반정밀 부동소수점 형식(torch.float16)으로 캐스팅하여 압축함으로써 통신 오버헤드를 줄입니다.
  • BF16 압축 훅: FP16과 유사하지만, 특정 하드웨어에서 더 효율적일 수 있는 Brain Floating Point 형식(torch.bfloat16)을 사용합니다.
  • PowerSGD Hook: 높은 압축률을 제공하여 대역폭 제한이 있는 분산 학습을 가속화할 수 있는 고급 그래디언트 압축 알고리즘입니다.

이 튜토리얼에서는 Accelerate에서 제공하는 유틸리티를 사용하여 DDP 통신 훅을 빠르게 설정하고 학습을 수행하는 방법을 확인합니다. 새 코드 한 줄만 추가하면 될 정도로 간단할 수 있습니다. 이는 Accelerate 라이브러리를 사용하여 분산 학습에서 DDP 통신 훅을 활용해 그라디언트 통신을 최적화하는 방법을 보여줍니다.

FP16 압축 훅

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 압축 훅

BF16 Compression Hook API는 실험적이며, 2.9.6보다 높은 NCCL 버전을 요구합니다.

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 훅

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 Communication Hooks 유틸리티

통신 훅을 통한 선택적 기능 지원을 위해 두 가지 추가 유틸리티가 있습니다.

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은 특정 communication hook에 필요한 추가 상태 정보를 전달할 수 있게 합니다. 이는 학습 단계 간에 하이퍼파라미터와 내부 상태를 유지해야 하는 PowerSGD와 같은 stateful hook에 특히 유용합니다. 아래는 PowerSGD hook과 함께 comm_state_option를 사용하는 예시입니다.

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