Accelerate documentation
Melakukan akumulasi gradien dengan Accelerate
Melakukan akumulasi gradien dengan Accelerate
Gradient accumulation adalah teknik di mana Anda dapat melatih model dengan ukuran batch yang lebih besar daripada yang biasanya mampu di muat ke dalam memori mesin Anda. Hal ini dilakukan dengan mengakumulasi gradien selama beberapa batch, dan hanya melakukan langkah optimizer setelah sejumlah batch tertentu telah diproses.
Secara teknis, kode akumulasi gradien standar akan berfungsi dengan baik dalam setup terdistribusi, tetapi ini bukan metode paling efisien untuk melakukannya dan Anda mungkin akan mengalami perlambatan yang cukup signifikan!
Pada tutorial ini, Anda akan melihat cara mengatur akumulasi gradien dengan cepat dan melaksanakannya menggunakan utilitas yang disediakan di Accelerate, yang pada akhirnya hanya membutuhkan penambahan satu baris kode baru!
Contoh ini akan menggunakan loop pelatihan PyTorch yang sangat sederhana yang melakukan akumulasi gradien setiap dua batch:
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()Mengonversinya ke Accelerate
Pertama, kode yang ditunjukkan sebelumnya akan dikonversi untuk memanfaatkan Accelerate tanpa helper akumulasi gradien khusus:
+ 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()Dalam kondisi saat ini, kode ini tidak akan melakukan akumulasi gradien secara efisien karena proses yang disebut sinkronisasi gradien. Baca lebih lanjut tentang hal tersebut di tutorial Konsep!
Membiarkan Accelerate menangani akumulasi gradien
Yang tersisa sekarang hanyalah membiarkan Accelerate menangani akumulasi gradien untuk kita. Untuk melakukannya, Anda harus meneruskan parameter gradient_accumulation_steps ke Accelerator, yang menentukan jumlah langkah yang harus dilakukan sebelum setiap pemanggilan step() dan bagaimana secara otomatis menyesuaikan loss selama pemanggilan backward():
from accelerate import Accelerator
- accelerator = Accelerator()
+ accelerator = Accelerator(gradient_accumulation_steps=2)Sebagai alternatif, Anda dapat meneruskan parameter gradient_accumulation_plugin ke __init__ objek Accelerator, yang memungkinkan Anda untuk lebih menyesuaikan perilaku akumulasi gradien.
Baca lebih lanjut tentang hal itu di dokumentasi GradientAccumulationPlugin.
Dari sini, Anda dapat menggunakan context manager accumulate() dari dalam loop pelatihan Anda untuk secara otomatis melakukan akumulasi gradien! Anda cukup membungkusnya di sekitar seluruh bagian pelatihan dari kode kami:
- for index, batch in enumerate(training_dataloader):
+ for batch in training_dataloader:
+ with accelerator.accumulate(model):
inputs, targets = batch
outputs = model(inputs)Anda dapat menghapus semua pemeriksaan khusus untuk nomor langkah dan penyesuaian loss:
- loss = loss / gradient_accumulation_steps
accelerator.backward(loss)
- if (index+1) % gradient_accumulation_steps == 0:
optimizer.step()
scheduler.step()
optimizer.zero_grad()Seperti yang dapat Anda lihat, Accelerator mampu melacak nomor batch yang sedang Anda proses dan secara otomatis mengetahui apakah perlu melakukan langkah pada optimizer yang telah disiapkan serta bagaimana cara menyesuaikan loss.
Biasanya, dengan akumulasi gradien, Anda perlu menyesuaikan jumlah langkah untuk mencerminkan perubahan total batch yang dilatih. Accelerate secara otomatis melakukan ini untuk Anda secara default. Di balik layar, kami menginstansiasi
GradientAccumulationPluginyang dikonfigurasi untuk melakukan hal ini.
state.GradientState disinkronkan dengan dataloader aktif yang sedang diiterasi. Dengan demikian, asumsi naifnya adalah bahwa ketika akhir dataloader tercapai, semuanya akan tersinkron dan sebuah langkah akan dilakukan. Untuk menonaktifkan ini, atur
sync_with_dataloadermenjadiFalsediGradientAccumulationPlugin:from accelerate import Accelerator from accelerate.utils import GradientAccumulationPlugin plugin = GradientAccumulationPlugin(sync_with_dataloader=False) accelerator = Accelerator(..., gradient_accumulation_plugin=plugin)
Kode yang selesai
Berikut adalah implementasi yang selesai untuk melakukan akumulasi gradien dengan 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()Penting bahwa hanya satu forward/backward yang dilakukan di dalam context manager
with accelerator.accumulate(model).
Untuk mempelajari lebih lanjut tentang apa yang dibungkus oleh keajaiban ini, baca panduan konsep Sinkronisasi Gradien
Contoh mandiri
Berikut adalah contoh mandiri yang dapat Anda jalankan untuk melihat akumulasi gradien dalam aksinya dengan 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.04000Akkumulasi gradien pada sampel pelatihan dengan ukuran variabel
Seperti yang ditunjukkan dalam blog-post ini, yang menyoroti kesalahan umum yang terjadi saat melakukan akumulasi gradien pada sampel pelatihan dengan ukuran yang bervariasi:
[…] untuk akumulasi gradien pada tugas tingkat token seperti pelatihan LM kausal, loss yang benar harus dihitung dengan membagi total loss di semua batch dalam satu langkah akumulasi gradien dengan total jumlah semua token non-padding di batch-batch tersebut. Ini tidak sama dengan rata-rata nilai loss per batch.
Dengan kata lain, beberapa penyesuaian harus dilakukan pada kerugian (losses) yang beroperasi pada tingkat token.
Hitung target yang benar-benar berkontribusi pada loss. Untuk model bahasa kausal Transformers yang menggeser label secara internal,
lalu lewati label yang belum digeser dan hitung (labels[..., 1:] != -100).sum(): label pertama tidak memiliki prediksi sebelumnya.
Jangan menggeser label tersebut lagi di data collator. Contoh akumulasi autoregresif menggunakan konvensi ini; fungsi loss lain mungkin menggunakan penargetan yang berbeda.
Kode kerangka
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()Contoh LM kausal yang mandiri
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()}")
Hasil pada satu perangkat - langkah akumulasi gradien diatur ke 1 dan batch_size diatur ke 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])Hasil pada konfigurasi dua perangkat - langkah akumulasi gradien diatur ke 2 dan batch_size diatur ke 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])Untuk melangkah lebih jauh:
Silakan lihat skrip contoh lengkap untuk pelatihan di dunia nyata di folder examples pada path accelerate/examples/by_feature/gradient_accumulation_for_autoregressive_models.py.
Menjalankannya pada beberapa konfigurasi pelatihan dengan ukuran batch global konstan yang sama dengan 32 menghasilkan grafik berikut:
Perhatikan bahwa kerugian pelatihan (training losses) sama persis hingga langkah pelatihan (training step) 20. Penyimpangan kecil setelah langkah pelatihan ini terjadi di akhir epoch pertama, karena, secara default, dataloader menduplikasi sampel di awal dataset ketika ukuran batch total (total batch size) tidak membagi dataset secara tepat.
Update on GitHub