8-bit Adam with Accelerator

You would like to customize your training loop with 8-bit Adam to reduce memory requirements of your model. Prepare the loop with 8-bit Adam for training.

Assume that an 8-bit Adam optimizer has been defined as adam_bnb_optim. Other training objects have been defined: model, train_dataloader, lr_scheduler, and accelerator.

本练习是课程的一部分

Efficient AI Model Training with PyTorch

查看课程

练习说明

Prepare the 8-bit Adam optimizer for distributed training.
Update the model parameters with the optimizer.
Zero the gradients with the optimizer.

交互式实操练习

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

# 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}")

编辑并运行代码