Gradient checkpointing avec Accelerator

Vous continuez à optimiser l’utilisation de la mémoire afin d’entraîner votre modèle de traduction automatique sur votre appareil. L’accumulation de gradients vous a permis d’entraîner efficacement avec des tailles de lot plus grandes. Poursuivez sur cette lancée en ajoutant le gradient checkpointing pour réduire l’empreinte mémoire de votre modèle.

Les objets model, train_dataloader et accelerator ont été préalablement définis.

Cet exercice fait partie du cours

Entraîner efficacement des modèles d’IA avec PyTorch

Afficher le cours

Instructions

Activez le gradient checkpointing sur model.
Configurez un gestionnaire de contexte Accelerator pour activer l’accumulation de gradients sur model.

Exercice interactif pratique

Essayez cet exercice en complétant cet exemple de code.

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

Modifier et exécuter le code