使用 Accelerator 进行梯度累积

您正在训练一个语言模型，通过改写复杂句子来简化翻译，但 GPU 内存不足。梯度累积可以让模型等效地在更大的批量上训练，同时实际使用能放入内存的小批量。您希望显式编写训练循环以便清楚其结构，因此使用 Accelerator。请注意，本练习实际在 CPU 上运行，但在 GPU 上代码相同。

model、train_dataloader、optimizer 和 lr_scheduler 都已预先定义。

本练习是课程的一部分

使用 PyTorch 高效训练 AI 模型

查看课程

练习说明

配置 Accelerator() 以使用 2 个步骤的梯度累积。
设置一个 Accelerator 上下文管理器，为 model 启用梯度累积。

交互式实操练习

通过完成这段示例代码来试试这个练习。

# Configure Accelerator
accelerator = ____(____=____)
model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(model, optimizer, train_dataloader, lr_scheduler)
for batch in train_dataloader:
    # Set up an Accelerator context manager
    with ____.____(____):
        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()
        print(f"Loss = {loss}")

编辑并运行代码