Local SGD dengan Accelerator
Anda telah mengimplementasikan akumulasi gradien dan checkpointing gradien untuk merapikan penggunaan memori pada model terjemahan bahasa Anda. Pelatihan masih agak lambat, jadi Anda memutuskan menambahkan local SGD ke loop pelatihan untuk meningkatkan efisiensi komunikasi antarperangkat. Bangun loop pelatihan dengan local SGD!
model, train_dataloader, dan accelerator telah didefinisikan sebelumnya, dan LocalSGD sudah diimpor.
Latihan ini merupakan bagian dari kursus
Pelatihan Model AI Efisien dengan PyTorch
Instruksi latihan
- Atur
local_sgd_stepsuntuk menyinkronkan gradien setiap delapan langkah. - Jalankan langkah pada manajer konteks local SGD.
Latihan interaktif langsung praktik
Cobalah latihan ini dengan melengkapi kode contoh ini.
# Set up a context manager to synchronize gradients every eight steps
with LocalSGD(accelerator=accelerator, model=model, local_sgd_steps=____, enabled=True) as local_sgd:
for batch in train_dataloader:
with accelerator.accumulate(model):
inputs, targets = batch["input_ids"], batch["labels"]
outputs = model(inputs, labels=targets)
loss = outputs.loss
accelerator.backward(loss)
optimizer.step()
lr_scheduler.step()
optimizer.zero_grad()
# Step the local SGD context manager
local_sgd.____()