使用 Accelerator 的 8-bit Adam

您希望自定义训练循环，使用 8-bit Adam 来降低模型的内存占用。请为训练准备基于 8-bit Adam 的循环。

假设已经定义了一个 8-bit Adam 优化器 adam_bnb_optim。其他训练对象也已定义：model、train_dataloader、lr_scheduler 和 accelerator。

本练习是课程的一部分

使用 PyTorch 高效训练 AI 模型

查看课程

练习说明

为分布式训练准备 8-bit Adam 优化器。
使用该优化器更新模型参数。
使用该优化器将梯度清零。

交互式实操练习

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

# Prepare the 8-bit Adam optimizer for distributed training
model, adam_bnb_optim, train_dataloader, lr_scheduler = accelerator.prepare(model, ____, train_dataloader, lr_scheduler)

for batch in train_dataloader:
    inputs, targets = batch["input_ids"], batch["labels"]
    outputs = model(inputs, labels=targets)
    loss = outputs.loss
    accelerator.backward(loss)
    # Update the model parameters
    adam_bnb_optim.____()
    lr_scheduler.step()
    # Zero the gradients
    adam_bnb_optim.____()
    print(f"Loss = {loss}")

编辑并运行代码