Local SGD กับ Accelerator
คุณได้ใช้งาน gradient accumulation และ gradient checkpointing เพื่อจัดการหน่วยความจำสำหรับโมเดลแปลภาษาแล้ว แต่การเทรนยังค่อนข้างช้าอยู่ จึงตัดสินใจเพิ่ม local SGD เข้าไปใน training loop เพื่อเพิ่มประสิทธิภาพการสื่อสารระหว่างอุปกรณ์ มาสร้าง training loop ที่มี local SGD กันเลย!
model, train_dataloader, และ accelerator ถูกกำหนดไว้ล่วงหน้าแล้ว และมีการ import LocalSGD เรียบร้อยแล้ว
แบบฝึกหัดนี้เป็นส่วนหนึ่งของหลักสูตร
การเทรน AI Model อย่างมีประสิทธิภาพด้วย PyTorch
คำแนะนำการฝึกหัด
- กำหนด
local_sgd_stepsให้ซิงโครไนซ์ gradient ทุก 8 สเต็ป - เรียกใช้เมธอด step ของ context manager local SGD
แบบฝึกหัดเชิงโต้ตอบแบบลงมือทำ
ลองทำแบบฝึกหัดนี้โดยเติมโค้ดตัวอย่างนี้ให้สมบูรณ์
# 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.____()