AdamW with Accelerator

You want to customize your training loop using Accelerator and use AdamW as a benchmark optimizer for your language translation model. Build the training loop to use AdamW.

Some training objects have been pre-loaded and defined, including model, train_dataloader, and accelerator.

This exercise is part of the course

Efficient AI Model Training with PyTorch

View Course

Exercise instructions

Prepare training objects for distributed training before the loop.
Update the model parameters in the training loop.

Hands-on interactive exercise

Have a go at this exercise by completing this sample code.

optimizer = AdamW(params=model.parameters())

# Prepare training objects for distributed training
model, optimizer, train_dataloader = ____.____(model, optimizer, train_dataloader)

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
    ____.____()
    optimizer.zero_grad()
    print(f"Loss = {loss}")

Edit and Run Code