Gradient checkpointing con Accelerator

Stai continuando a ottimizzare l'uso della memoria per poter addestrare sul tuo dispositivo il tuo modello di traduzione automatica. La gradient accumulation ti ha aiutato ad addestrare in modo efficace con batch più grandi. Prosegui su questa linea aggiungendo il gradient checkpointing per ridurre l'impronta in memoria del modello.

model, train_dataloader e accelerator sono già stati definiti.

Questo esercizio fa parte del corso

Efficient AI Model Training with PyTorch

Visualizza corso

Istruzioni dell'esercizio

Abilita il gradient checkpointing su model.
Imposta un context manager di Accelerator per abilitare la gradient accumulation su model.

esercizio interattivo pratico

Prova questo esercizio completando questo codice di esempio.

# Enable gradient checkpointing on the model
____.____()

for batch in train_dataloader:
    with accelerator.accumulate(model):
        inputs, targets = batch["input_ids"], batch["labels"]
        # Get the outputs from a forward pass of the model
        ____ = ____(____, labels=targets)
        loss = outputs.loss
        accelerator.backward(loss)
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()
        print(f"Loss = {loss}")

Modifica ed esegui il codice