Local SGD với Accelerator
Bạn đã triển khai cộng dồn gradient (gradient accumulation) và gradient checkpointing để tối ưu việc sử dụng bộ nhớ cho mô hình dịch ngôn ngữ. Quá trình huấn luyện vẫn hơi chậm, nên bạn quyết định thêm local SGD vào vòng lặp huấn luyện để cải thiện hiệu quả giao tiếp giữa các thiết bị. Hãy xây dựng vòng lặp huấn luyện với local SGD!
model, train_dataloader và accelerator đã được định nghĩa sẵn, và LocalSGD đã được import.
Bài tập này là một phần của khóa học
Huấn luyện Mô hình AI Hiệu quả với PyTorch
Hướng dẫn bài tập
- Đặt
local_sgd_stepsđể đồng bộ hóa gradient sau mỗi tám bước. - Gọi bước cho trình quản lý ngữ cảnh local SGD.
Bài tập tương tác thực hành trực tiếp
Hãy thử làm bài tập này bằng cách hoàn thành đoạn mã mẫu này.
# 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.____()