Point de contrôle

Lors de l’entraînement d’un modèle PyTorch avec Accelerate, vous souhaitez souvent sauvegarder et reprendre un état d’entraînement. Pour ce faire, il est nécessaire de sauvegarder et de charger le modèle, l’optimiseur, les générateurs de nombres aléatoires (RNG) et le GradScaler. Accelerate propose deux fonctions pratiques pour réaliser cela rapidement :

• Utilisez save_state() pour enregistrer tout ce qui précède dans un emplacement de dossier
• Utilisez load_state() pour charger tout ce qui a été enregistré lors d'un save_state précédent

Pour personnaliser davantage l’emplacement et la manière dont les états sont enregistrés via save_state(), la classe ProjectConfiguration peut être utilisée. Par exemple, si automatic_checkpoint_naming est activé, chaque point de contrôle enregistré sera ensuite situé à Accelerator.project_dir/checkpoints/checkpoint_{checkpoint_number}.

Il convient de noter que l'on s'attend à ce que ces états proviennent du même script d'entraînement ; ils ne doivent pas provenir de deux scripts distincts.

• En utilisant register_for_checkpointing(), vous pouvez enregistrer des objets personnalisés pour qu’ils soient automatiquement stockés ou chargés par les deux fonctions précédentes, à condition que l’objet dispose d’une fonctionnalité state_dict et d’une fonctionnalité load_state_dict. Cela peut inclure des objets tels qu’un planificateur de taux d’apprentissage.

Voici un exemple succinct utilisant le checkpointing pour enregistrer et recharger un état pendant l’entraînement :

from accelerate import Accelerator
import torch

accelerator = Accelerator(project_dir="my/save/path")

my_scheduler = torch.optim.lr_scheduler.StepLR(my_optimizer, step_size=1, gamma=0.99)
my_model, my_optimizer, my_training_dataloader = accelerator.prepare(my_model, my_optimizer, my_training_dataloader)

# Register the LR scheduler
accelerator.register_for_checkpointing(my_scheduler)

# Save the starting state
accelerator.save_state()

device = accelerator.device
my_model.to(device)

# Perform training
for epoch in range(num_epochs):
    for batch in my_training_dataloader:
        my_optimizer.zero_grad()
        inputs, targets = batch
        inputs = inputs.to(device)
        targets = targets.to(device)
        outputs = my_model(inputs)
        loss = my_loss_function(outputs, targets)
        accelerator.backward(loss)
        my_optimizer.step()
    my_scheduler.step()

# Restore the previous state
accelerator.load_state("my/save/path/checkpointing/checkpoint_0")

Restaurer l'état du DataLoader

Après la reprise depuis un point de contrôle, il peut également être souhaitable de reprendre à un point précis du DataLoader actif si l’état a été enregistré en cours d’époque. Vous pouvez utiliser skip_first_batches() pour ce faire.

from accelerate import Accelerator

accelerator = Accelerator(project_dir="my/save/path")

train_dataloader = accelerator.prepare(train_dataloader)
accelerator.load_state("my_state")

# Assume the checkpoint was saved 100 steps into the epoch
skipped_dataloader = accelerator.skip_first_batches(train_dataloader, 100)

# After the first iteration, go back to `train_dataloader`

# First epoch
for batch in skipped_dataloader:
    # Do something
    pass

# Second epoch
for batch in train_dataloader:
    # Do something
    pass