Gradient checkpointing với Accelerator
Bạn đang tiếp tục tối ưu hóa việc sử dụng bộ nhớ để có thể huấn luyện mô hình dịch ngôn ngữ ngay trên thiết bị của mình. Gradient accumulation đã giúp bạn huấn luyện hiệu quả với kích thước batch lớn hơn. Hãy phát huy tiếp công việc này bằng cách thêm gradient checkpointing để giảm dấu chân bộ nhớ của mô hình.
model, train_dataloader và accelerator đã được định nghĩa sẵn.
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
- Bật gradient checkpointing trên
model. - Thiết lập một context manager của
Acceleratorđể bật gradient accumulation trênmodel.
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.
# Enable gradient checkpointing on the model
____.____()
for batch in train_dataloader:
with accelerator.accumulate(model):
inputs, targets = batch["input_ids"], batch["labels"]
# Get the outputs from a forward pass of the model
____ = ____(____, labels=targets)
loss = outputs.loss
accelerator.backward(loss)
optimizer.step()
lr_scheduler.step()
optimizer.zero_grad()
print(f"Loss = {loss}")