Acumulação de gradientes com Accelerator

Você está treinando um modelo de linguagem para simplificar traduções ao parafrasear sentenças complexas, mas sua GPU está ficando sem memória. A acumulação de gradientes permite que o modelo treine efetivamente em lotes maiores usando mini-lotes que cabem na memória. Você prefere escrever explicitamente o loop de treino para visualizar sua estrutura, então está usando o Accelerator. Observe que este exercício roda no CPU, mas o código é o mesmo para a GPU.

O model, train_dataloader, optimizer e lr_scheduler já foram definidos.

Este exercício faz parte do curso

Treinamento Eficiente de Modelos de IA com PyTorch

Ver curso

Instruções do exercício

Configure o Accelerator() para usar acumulação de gradientes com duas etapas.
Configure um gerenciador de contexto do Accelerator para habilitar a acumulação de gradientes para o model.

Exercício interativo prático

Experimente este exercício completando este código de exemplo.

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

Editar e executar o código