SGD local avec Accelerator

Vous avez mis en place l'accumulation de gradients et le gradient checkpointing pour optimiser l'utilisation de la mémoire de votre modèle de traduction automatique. L'entraînement reste un peu lent, alors vous décidez d'ajouter le SGD local à votre boucle d'entraînement afin d'améliorer l'efficacité des communications entre appareils. Créez la boucle d'entraînement avec le SGD local!

Les objets model, train_dataloader et accelerator sont déjà définis, et LocalSGD a été importé.

Cette activité fait partie du cours

Entraînement efficace de modèles d’IA avec PyTorch

Voir le cours

Instructions de l’exercice

Définissez local_sgd_steps pour synchroniser les gradients toutes les huit étapes.
Faites avancer le gestionnaire de contexte du SGD local.

Exercice interactif pratique

Essayez cet exercice en complétant ce code d’exemple.

# Set up a context manager to synchronize gradients every eight steps
with LocalSGD(accelerator=accelerator, model=model, local_sgd_steps=____, enabled=True) as local_sgd:
    for batch in train_dataloader:
        with accelerator.accumulate(model):
            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()
            # Step the local SGD context manager
            local_sgd.____()

Modifier et exécuter le code