Accelerate documentation
Accelerate를 사용한 그라디언트 누적 수행
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).
Accelerate를 사용한 그라디언트 누적 수행
Gradient accumulation은 기계의 메모리에 일반적으로 담을 수 없는 더 큰 batch size로 학습할 수 있는 기법입니다. 여러 batch에 걸쳐 gradient를 누적하고, 특정 수의 batch가 처리된 후에만 optimizer를 step하는 방식으로 동작합니다.
기술적으로는 표준 gradient accumulation 코드가 분산 환경에서도 정상적으로 작동하지만, 가장 효율적인 방법은 아니며 상당한 성능 저하가 발생할 수 있습니다!
이 튜토리얼에서는 Accelerate에서 제공하는 유틸리티를 사용하여 그라디언트 어큐뮬레이션을 빠르게 설정하고 실행하는 방법을 확인합니다. 이를 위해 추가해야 하는 코드는 단 한 줄뿐입니다!
이 예제는 두 개의 배치마다 그라디언트 축적을 수행하는 매우 단순한 PyTorch 학습 루프를 사용합니다:
device = "cuda"
model.to(device)
gradient_accumulation_steps = 2
for index, batch in enumerate(training_dataloader):
inputs, targets = batch
inputs = inputs.to(device)
targets = targets.to(device)
outputs = model(inputs)
loss = loss_function(outputs, targets)
loss = loss / gradient_accumulation_steps
loss.backward()
if (index + 1) % gradient_accumulation_steps == 0:
optimizer.step()
scheduler.step()
optimizer.zero_grad()Accelerate로 변환
앞서 보여준 코드를 특별한 그라디언트 축적 헬퍼 없이 Accelerate를 활용하도록 변환합니다:
+ from accelerate import Accelerator
+ accelerator = Accelerator()
+ model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+ model, optimizer, training_dataloader, scheduler
+ )
for index, batch in enumerate(training_dataloader):
inputs, targets = batch
- inputs = inputs.to(device)
- targets = targets.to(device)
outputs = model(inputs)
loss = loss_function(outputs, targets)
loss = loss / gradient_accumulation_steps
+ accelerator.backward(loss)
if (index+1) % gradient_accumulation_steps == 0:
optimizer.step()
scheduler.step()
optimizer.zero_grad()현재 상태에서는 개념 튜토리얼에서 자세히 설명하는 그래디언트 동기화 과정 때문에 이 코드는 그래디언트 축적을 효율적으로 수행하지 못합니다.
Accelerate로 그래디언트 축적을 처리
이제 남은 작업은 Accelerate가 그라디언트 축적을 처리하도록 하는 것입니다. 이를 위해 Accelerator에 gradient_accumulation_steps 매개변수를 전달하여, step() 호출 전 수행할 스텝 수와 backward() 호출 시 손실을 자동으로 조정하는 방법을 지정해야 합니다:
from accelerate import Accelerator
- accelerator = Accelerator()
+ accelerator = Accelerator(gradient_accumulation_steps=2)대신 Accelerator 객체의 __init__에 gradient_accumulation_plugin 매개변수를 전달하여 그래디언트 축적 동작을 더 자세히 사용자 정의할 수 있습니다.
이에 대한 자세한 내용은 GradientAccumulationPlugin 문서를 참고하세요.
여기서부터는 학습 루프 내부에서 accumulate() 컨텍스트 매니저를 사용하여 그라디언트 축적을 자동으로 수행할 수 있습니다! 코드 전체의 학습 부분을 이 컨텍스트 매니저로 감싸기만 하면 됩니다:
- for index, batch in enumerate(training_dataloader):
+ for batch in training_dataloader:
+ with accelerator.accumulate(model):
inputs, targets = batch
outputs = model(inputs)step 번호와 손실 조정과 관련된 모든 특수한 체크를 제거할 수 있습니다:
- loss = loss / gradient_accumulation_steps
accelerator.backward(loss)
- if (index+1) % gradient_accumulation_steps == 0:
optimizer.step()
scheduler.step()
optimizer.zero_grad()보시는 바와 같이 Accelerator는 현재 배치 번호를 추적할 수 있으며, 준비된 옵티마이저를 단계별로 진행할지 여부와 손실을 어떻게 조정할지를 자동으로 파악합니다.
일반적으로 gradient accumulation을 사용할 경우, 학습에 사용하는 총 배치 수의 변화를 반영하도록 스텝 수를 조정해야 합니다. Accelerate는 기본적으로 이를 자동으로 처리합니다. 내부적으로는 이를 수행하도록 구성된
GradientAccumulationPlugin을 인스턴스화합니다.
state.GradientState는 현재 반복 중인 활성 dataloader와 동기화됩니다. 따라서 dataloader의 끝까지 도달하면 모든 것이 동기화되고 step이 수행된다고 단순하게 가정합니다. 이를 비활성화하려면
GradientAccumulationPlugin에서sync_with_dataloader를False로 설정하세요:from accelerate import Accelerator from accelerate.utils import GradientAccumulationPlugin plugin = GradientAccumulationPlugin(sync_with_dataloader=False) accelerator = Accelerator(..., gradient_accumulation_plugin=plugin)
완성된 코드
Accelerate를 사용하여 그라디언트 축적을 수행하는 완성된 구현은 다음과 같습니다.
from accelerate import Accelerator
accelerator = Accelerator(gradient_accumulation_steps=2)
model, optimizer, training_dataloader, scheduler = accelerator.prepare(
model, optimizer, training_dataloader, scheduler
)
for batch in training_dataloader:
with accelerator.accumulate(model):
inputs, targets = batch
outputs = model(inputs)
loss = loss_function(outputs, targets)
accelerator.backward(loss)
optimizer.step()
scheduler.step()
optimizer.zero_grad()컨텍스트 매니저
with accelerator.accumulate(model)안에서는 forward/backward가 한 번만 수행되어야 합니다.
이 마법이 무엇을 감싸고 있는지 자세히 알아보려면 Gradient Synchronization 개념 가이드를 읽어보세요.
독립형 예제
Accelerate에서 그라디언트 축적이 어떻게 작동하는지 확인하기 위해 실행할 수 있는 독립형 예제를 소개합니다:
import torch
import copy
from accelerate import Accelerator
from accelerate.utils import set_seed
from torch.utils.data import TensorDataset, DataLoader
# seed
set_seed(0)
# define toy inputs and labels
x = torch.tensor([1., 2., 3., 4., 5., 6., 7., 8.])
y = torch.tensor([2., 4., 6., 8., 10., 12., 14., 16.])
gradient_accumulation_steps = 4
per_device_batch_size = len(x) // gradient_accumulation_steps
# define dataset and dataloader
dataset = TensorDataset(x, y)
dataloader = DataLoader(dataset, batch_size=per_device_batch_size)
# define model, optimizer and loss function
class SimpleLinearModel(torch.nn.Module):
def __init__(self):
super(SimpleLinearModel, self).__init__()
self.weight = torch.nn.Parameter(torch.zeros((1, 1)))
def forward(self, inputs):
return inputs @ self.weight
model = SimpleLinearModel()
model_clone = copy.deepcopy(model)
criterion = torch.nn.MSELoss()
model_optimizer = torch.optim.SGD(model.parameters(), lr=0.02)
accelerator = Accelerator(gradient_accumulation_steps=gradient_accumulation_steps)
model, model_optimizer, dataloader = accelerator.prepare(model, model_optimizer, dataloader)
model_clone_optimizer = torch.optim.SGD(model_clone.parameters(), lr=0.02)
print(f"initial model weight is {model.weight.mean().item():.5f}")
print(f"initial model weight is {model_clone.weight.mean().item():.5f}")
for i, (inputs, labels) in enumerate(dataloader):
with accelerator.accumulate(model):
inputs = inputs.view(-1, 1)
print(i, inputs.flatten())
labels = labels.view(-1, 1)
outputs = model(inputs)
loss = criterion(outputs, labels)
accelerator.backward(loss)
model_optimizer.step()
model_optimizer.zero_grad()
loss = criterion(x.view(-1, 1) @ model_clone.weight, y.view(-1, 1))
model_clone_optimizer.zero_grad()
loss.backward()
model_clone_optimizer.step()
print(f"w/ accumulation, the final model weight is {model.weight.mean().item():.5f}")
print(f"w/o accumulation, the final model weight is {model_clone.weight.mean().item():.5f}")initial model weight is 0.00000
initial model weight is 0.00000
0 tensor([1., 2.])
1 tensor([3., 4.])
2 tensor([5., 6.])
3 tensor([7., 8.])
w/ accumulation, the final model weight is 2.04000
w/o accumulation, the final model weight is 2.04000크기가 다양한 학습 샘플에 대한 Gradient accumulation
이 블로그 포스트에서 지적했듯이, 가변 크기의 학습 샘플에서 그라디언트 축적을 수행할 때 발생하는 일반적인 오류를 지적합니다:
[…] 토큰 수준 작업(예: causal LM 학습)에 대한 gradient accumulation의 경우, 올바른 loss는 gradient accumulation 단계 내 모든 배치에 걸친 총 loss를 해당 배치들의 모든 non padding 토큰의 총 개수로 나누어 계산해야 합니다. 이는 배치별 loss 값의 평균과 동일하지 않습니다.
즉, 토큰 수준에서 동작하는 손실 함수에 대해 일부 조치가 필요합니다.
실제로 손실에 기여하는 타깃의 개수를 세어 보세요. 레이블을 내부적으로 시프트하는 Transformers의 causal language model의 경우, 시프트되지 않은 레이블을 전달하고 (labels[..., 1:] != -100).sum()을 계산하세요. 첫 번째 레이블에는 이전 예측이 없기 때문입니다. 데이터 콜레이터에서 해당 레이블을 다시 시프트하지 마세요. autoregressive accumulation 예제는 이 관례를 사용합니다. 다른 손실 함수는 다른 타깃 정렬 방식을 사용할 수 있습니다.
스켈레톤 코드
from accelerate import Accelerator
import math
import contextlib
gradient_accumulation_steps = 2
accelerator = Accelerator(gradient_accumulation_steps=gradient_accumulation_steps)
model, optimizer, training_dataloader, scheduler = accelerator.prepare(
model, optimizer, training_dataloader, scheduler
)
training_iterator = iter(training_dataloader)
num_samples_in_epoch = len(training_dataloader)
remainder = num_samples_in_epoch % gradient_accumulation_steps
remainder = remainder if remainder != 0 else gradient_accumulation_steps
total_updates = math.ceil(num_samples_in_epoch / gradient_accumulation_steps)
total_batched_samples = 0
for update_step in range(total_updates):
# In order to correctly the total number of non-padded tokens on which we'll compute the cross-entropy loss
# we need to pre-load the full local batch - i.e the next per_device_batch_size * accumulation_steps samples
batch_samples = []
num_batches_in_step = gradient_accumulation_steps if update_step != (total_updates - 1) else remainder
for _ in range(num_batches_in_step):
batch_samples += [next(training_iterator)]
# get local num items in batch
num_items_in_batch = sum([(batch["labels"].ne(-100)).sum() for batch in batch_samples])
# to compute it correctly in a multi-device DDP training, we need to gather the total number of items in the full batch.
num_items_in_batch = accelerator.gather(num_items_in_batch).sum().item()
for i, batch in enumerate(batch_samples):
# if we perform gradient accumulation in a multi-devices set-up, we want to avoid unnecessary communications when accumulating
# cf: https://muellerzr.github.io/blog/gradient_accumulation.html
if (i < len(batch_samples) - 1 and accelerator.num_processes > 1):
ctx = model.no_sync
else:
ctx = contextlib.nullcontext
total_batched_samples += 1
with ctx():
inputs, targets = batch
outputs = model(inputs)
loss = loss_function(outputs, targets) # the loss function should sum over samples rather than averaging
# We multiply by num_processes because the DDP calculates the average gradient across all devices whereas dividing by num_items_in_batch already takes into account all devices
# Same reason for gradient_accumulation_steps, but this times it's Accelerate that calculate the average gradient across the accumulated steps
loss = (loss * gradient_accumulation_steps * accelerator.num_processes) / num_items_in_batch
accelerator.backward(loss)
# Sync gradients and perform optimization steps once every gradient_accumulation_steps
optimizer.step()
scheduler.step()
optimizer.zero_grad()자체 완결형 causal LM 예제
import torch
import copy
from accelerate import Accelerator
from accelerate.utils import set_seed
from accelerate.logging import get_logger
from torch.utils.data import Dataset, DataLoader
import math
import contexlib
# seed
set_seed(0)
logger = get_logger(__name__)
class MyDataset(Dataset):
def __init__(self, num_samples):
super().__init__()
self.len = num_samples
def __getitem__(self, index):
input_ids = torch.arange(1, index+2, dtype=torch.float32)
labels = torch.remainder(input_ids, 2)
return {"input_ids": input_ids, "labels": labels}
def __len__(self):
return self.len
def collate_fn(features):
input_ids = torch.nn.utils.rnn.pad_sequence([f["input_ids"] for f in features], batch_first=True, padding_value=-100)
labels = torch.nn.utils.rnn.pad_sequence([f["labels"] for f in features], batch_first=True, padding_value=-100)
return {"input_ids": input_ids[..., None], "labels": labels[..., None]}
# define toy inputs and labels
gradient_accumulation_steps = 2
per_device_batch_size = 4
# define accelerator
accelerator = Accelerator(gradient_accumulation_steps=gradient_accumulation_steps)
# define dataset and dataloader
# for this toy example, we'll compute gradient descent over one single global batch
dataset = MyDataset(per_device_batch_size*gradient_accumulation_steps*accelerator.num_processes)
dataloader = DataLoader(dataset, batch_size=per_device_batch_size, collate_fn=collate_fn)
# define model, model_optimizer and loss function
model = torch.nn.Linear(1, 2, bias=False)
model_clone = copy.deepcopy(model)
criterion = torch.nn.CrossEntropyLoss(reduction="sum") # must sum over samples rather than averaging
model_optimizer = torch.optim.SGD(model.parameters(), lr=0.08)
logger.warning(f"initial model weight is {model.weight.detach().cpu().squeeze()}")
logger.warning(f"initial model clone weight is {model_clone.weight.detach().cpu().squeeze()}")
# prepare artifacts - accelerator handles device placement and dataloader splitting
model, model_optimizer = accelerator.prepare(model, model_optimizer)
dataloader = accelerator.prepare_data_loader(dataloader, device_placement=True)
training_iterator = iter(dataloader)
num_samples_in_epoch = len(dataloader)
remainder = num_samples_in_epoch % gradient_accumulation_steps
remainder = remainder if remainder != 0 else gradient_accumulation_steps
total_gradient_updates = math.ceil(num_samples_in_epoch / gradient_accumulation_steps)
total_batched_samples = 0
for update_step in range(total_gradient_updates):
# In order to correctly the total number of non-padded tokens on which we'll compute the cross-entropy loss
# we need to pre-load the full local batch - i.e the next per_device_batch_size * accumulation_steps samples
batch_samples = []
num_batches_in_step = gradient_accumulation_steps if update_step != (total_gradient_updates - 1) else remainder
for _ in range(num_batches_in_step):
batch_samples += [next(training_iterator)]
# get local num items in batch
local_num_items_in_batch = sum([(batch["labels"].ne(-100)).sum() for batch in batch_samples])
logger.warning(f"Step {update_step} - Device {accelerator.process_index} - num items in the local batch {local_num_items_in_batch}", main_process_only=False)
# to compute it correctly in a multi-device DDP training, we need to gather the total number of items in the full batch.
num_items_in_batch = accelerator.gather(local_num_items_in_batch).sum().item()
logger.warning(f"Total num items {num_items_in_batch}")
for i, batch in enumerate(batch_samples):
inputs, labels = batch["input_ids"], batch["labels"]
total_batched_samples += 1
# if we perform gradient accumulation in a multi-devices set-up, we want to avoid unnecessary communications when accumulating
# cf: https://muellerzr.github.io/blog/gradient_accumulation.html
if (i < len(batch_samples) - 1 and accelerator.num_processes > 1):
ctx = model.no_sync
else:
ctx = contextlib.nullcontext
with ctx():
outputs = model(inputs)
loss = criterion(outputs.view(-1, 2), labels.view(-1).to(torch.int64))
# We multiply by num_processes because the DDP calculates the average gradient across all devices whereas dividing by num_items_in_batch already takes into account all devices
# Same reason for gradient_accumulation_steps, but this times it's Accelerate that calculate the average gradient across the accumulated steps
loss = (loss * gradient_accumulation_steps * accelerator.num_processes) / num_items_in_batch
accelerator.backward(loss)
model_optimizer.step()
model_optimizer.zero_grad()
logger.warning(f"Device {accelerator.process_index} - w/ accumulation, the final model weight is {accelerator.unwrap_model(model).weight.detach().cpu().squeeze()}", main_process_only=False)
# We know do the same operation but on a single device and without gradient accumulation
if accelerator.is_main_process:
# prepare one single entire batch
dataloader = DataLoader(dataset, batch_size=len(dataset), collate_fn=collate_fn)
full_batch_without_accum = next(iter(dataloader))
total_inputs, total_labels = full_batch_without_accum["input_ids"], full_batch_without_accum["labels"]
model_clone_optimizer = torch.optim.SGD(model_clone.parameters(), lr=0.08)
# train the cloned model
loss = torch.nn.CrossEntropyLoss(reduction="mean")(model_clone(total_inputs).view(-1, 2), total_labels.view(-1).to(torch.int64))
model_clone_optimizer.zero_grad()
loss.backward()
model_clone_optimizer.step()
# We should have the same final weights.
logger.warning(f"w/o accumulation, the final model weight is {model_clone.weight.detach().cpu().squeeze()}")
gradient accumulation steps를 1로, batch_size를 8로 설정한 단일 디바이스에서의 결과:
initial model weight is tensor([-0.0075, 0.5364])
initial model clone weight is tensor([-0.0075, 0.5364])
Step 0 - Device 0 - num items in the local batch 36
Total num items 36
Device 0 - w/ accumulation, the final model weight is tensor([0.0953, 0.4337])
w/o accumulation, the final model weight is tensor([0.0953, 0.4337])2개의 장치로 구성된 환경에서의 결과 - gradient accumulation steps는 2로, batch_size는 4로 설정했습니다.
initial model weight is tensor([-0.0075, 0.5364])
initial model clone weight is tensor([-0.0075, 0.5364])
Step 0 - Device 0 - num items in the local batch 52
Step 0 - Device 1 - num items in the local batch 84
Total num items 136
Device 1 - w/ accumulation, the final model weight is tensor([0.2117, 0.3172])
Device 0 - w/ accumulation, the final model weight is tensor([0.2117, 0.3172])
w/o accumulation, the final model weight is tensor([0.2117, 0.3172])더 깊이 파고들기:
실제 학습 실행에 대한 완전한 예제 스크립트는 accelerate/examples/by_feature/gradient_accumulation_for_autoregressive_models.py 경로에 있는 examples 폴더에서 확인하실 수 있습니다.
global batch size가 32로 고정된 여러 학습 구성으로 실행하면 다음 그래프가 생성됩니다:
학습 단계 20까지의 학습 손실이 정확히 동일하다는 점에 유의하세요. 이 학습 단계 이후의 작은 편차는 첫 번째 에포크의 마지막 부분에 발생하며, 이는 기본적으로 전체 배치 크기가 데이터셋을 정확히 나누지 못할 때 데이터로더가 데이터셋의 처음에 있는 샘플을 복제하기 때문입니다.
Update on GitHub