Ajouter Accelerate à votre code

Chaque framework d’entraînement distribué a sa propre approche, ce qui peut nécessiter d’écrire beaucoup de code sur mesure pour l’adapter à votre code d’entraînement PyTorch et à votre environnement d’entraînement. Accelerate propose une manière conviviale d’interagir avec ces frameworks d’entraînement distribué sans avoir à apprendre les détails spécifiques de chacun. Accelerate s’occupe de ces détails pour vous, afin que vous puissiez vous concentrer sur le code d’entraînement et le mettre à l’échelle dans n’importe quel environnement d’entraînement distribué.

Dans ce tutoriel, vous apprendrez à adapter votre code PyTorch existant avec Accelerate pour vous mettre en route vers l’entraînement sur des systèmes distribués en toute simplicité ! Vous commencerez par une boucle d’entraînement PyTorch de base (il est supposé que tous les objets d’entraînement comme model et optimizer ont déjà été configurés) et vous intégrerez progressivement Accelerate.

device = "cuda"
model.to(device)

for batch in training_dataloader:
    optimizer.zero_grad()
    inputs, targets = batch
    inputs = inputs.to(device)
    targets = targets.to(device)
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
    loss.backward()
    optimizer.step()
    scheduler.step()

Accélérateur

La classe Accelerator est la principale pour adapter votre code afin qu’il fonctionne avec Accelerate. Elle connaît la configuration distribuée que vous utilisez, comme le nombre de processus distincts et le type de matériel. Cette classe donne également accès à de nombreuses méthodes nécessaires pour permettre à votre code PyTorch de fonctionner dans n’importe quel environnement d’entraînement distribué et pour gérer et exécuter les processus sur les différents appareils.

C’est pourquoi vous devez toujours commencer par importer et créer une instance Accelerator dans votre script.

from accelerate import Accelerator

accelerator = Accelerator()

Le Accelerator sait également sur quel appareil déplacer vos objets PyTorch, il est donc recommandé de laisser Accelerate s’en charger.

- device = "cuda"
+ device = accelerator.device
  model.to(device)

Préparer les objets PyTorch

Ensuite, vous devez préparer vos objets PyTorch (modèle, optimiseur, planificateur, etc.) pour l’entraînement distribué. La méthode prepare() s’occupe de placer votre modèle dans le conteneur approprié (comme un GPU unique ou multi-GPU) pour votre configuration d’entraînement, d’adapter l’optimiseur et le planificateur pour utiliser AcceleratedOptimizer et AcceleratedScheduler d’Accelerate, et de créer un nouveau dataloader qui peut être fractionné entre les processus.

Accelerate ne prépare que les objets qui héritent de leurs classes PyTorch respectives, telles que torch.optim.Optimizer.

Les objets PyTorch sont renvoyés dans le même ordre que celui dans lequel ils sont envoyés.

model, optimizer, training_dataloader, scheduler = accelerator.prepare(
    model, optimizer, training_dataloader, scheduler
)

Boucle d'entraînement

Enfin, supprimez les appels to(device) aux entrées et aux cibles dans la boucle d’entraînement, car les classes DataLoader d’Accelerate les placent automatiquement sur le bon appareil. Vous devez également remplacer le passage backward() habituel par la méthode backward() d’Accelerate, qui met à l’échelle les gradients pour vous et utilise la méthode backward() appropriée en fonction de votre configuration distribuée (par exemple, DeepSpeed ou Megatron).

-   inputs = inputs.to(device)
-   targets = targets.to(device)
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
-   loss.backward()
+   accelerator.backward(loss)

Assemblez le tout et votre nouvelle boucle d’entraînement Accelerate devrait maintenant ressembler à ceci !

from accelerate import Accelerator
accelerator = Accelerator()

device = accelerator.device
model, optimizer, training_dataloader, scheduler = accelerator.prepare(
    model, optimizer, training_dataloader, scheduler
)

for batch in training_dataloader:
    optimizer.zero_grad()
    inputs, targets = batch
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
    accelerator.backward(loss)
    optimizer.step()
    scheduler.step()

Fonctionnalités d'entraînement

Accelerate propose des fonctionnalités supplémentaires – comme l’accumulation de gradients, le clipping de gradients et l’entraînement en précision mixte, entre autres – que vous pouvez ajouter à votre script pour améliorer votre session d’entraînement. Explorons ces trois fonctionnalités.

Accumulation du gradient

L'accumulation de gradients vous permet d'entraîner sur des tailles de lot plus grandes en accumulant les gradients sur plusieurs lots avant de mettre à jour les poids. Cela peut être utile pour contourner les limitations de mémoire. Pour activer cette fonctionnalité dans Accelerate, spécifiez le paramètre gradient_accumulation_steps dans la classe Accelerator et ajoutez le gestionnaire de contexte accumulate() à votre script.

+ accelerator = Accelerator(gradient_accumulation_steps=2)
  model, optimizer, training_dataloader = accelerator.prepare(model, optimizer, training_dataloader)

  for input, label in training_dataloader:
+     with accelerator.accumulate(model):
          predictions = model(input)
          loss = loss_function(predictions, label)
          accelerator.backward(loss)
          optimizer.step()
          scheduler.step()
          optimizer.zero_grad()

Régularisation du gradient

La coupure du gradient est une technique visant à prévenir les « gradients explosifs », et Accelerate propose :

• clipgrad_value() pour limiter les gradients à une valeur minimale et maximale
• clipgrad_norm() pour normaliser les gradients à une valeur donnée

Précision mixte

La précision mixte accélère l’entraînement en utilisant un type de données de précision inférieure, comme fp16 (précision demi), pour calculer les gradients. Pour obtenir les meilleures performances avec Accelerate, la perte doit être calculée à l’intérieur de votre modèle (comme dans les modèles Transformers), car les calculs effectués en dehors du modèle sont réalisés en pleine précision.

Définissez le type de précision mixte à utiliser dans Accelerator, puis utilisez le gestionnaire de contexte autocast() pour convertir automatiquement les valeurs dans le type de données spécifié.

Accelerate permet la précision mixte automatique, de sorte que autocast() n'est nécessaire que s'il existe d'autres opérations de précision mixte en plus de celles effectuées sur la perte par backward(), qui gère déjà le mise à l'échelle.

+ accelerator = Accelerator(mixed_precision="fp16")
+ with accelerator.autocast():
      loss = complex_loss_function(outputs, target)

Enregistrer et charger

Accelerate peut également enregistrer et charger un modèle une fois l’entraînement terminé, ou bien enregistrer l’état du modèle et de l’optimiseur state, ce qui peut s’avérer utile pour reprendre l’entraînement.

Modèle

Une fois tous les processus terminés, déballer le modèle avec la méthode unwrap_model() avant de l’enregistrer, car la méthode prepare() a encapsulé votre modèle dans l’interface appropriée pour l’entraînement distribué. Si vous ne déballez pas le modèle, l’enregistrement du dictionnaire d’état du modèle enregistre également d’éventuelles couches supplémentaires issues du modèle plus grand, et vous ne pourrez pas recharger les poids dans votre modèle de base.

Vous devez utiliser la méthode save_model() pour déballer et enregistrer le dictionnaire d'état du modèle. Cette méthode peut également enregistrer un modèle dans des points de contrôle fractionnés ou au format safetensors.

accelerator.wait_for_everyone()
accelerator.save_model(model, save_directory)

unwrapped_model = accelerator.unwrap_model(model)
path_to_checkpoint = os.path.join(save_directory,"pytorch_model.bin")
unwrapped_model.load_state_dict(torch.load(path_to_checkpoint))

accelerator.wait_for_everyone()
accelerator.save_model(model, save_directory, max_shard_size="1GB", safe_serialization=True)

load_checkpoint_in_model(unwrapped_model, save_directory, device_map={"":device})

État

Lors de l'entraînement, vous pouvez vouloir enregistrer l'état actuel du modèle, de l'optimiseur, des générateurs aléatoires et, le cas échéant, des planificateurs de taux d'apprentissage, afin de pouvoir les restaurer dans le même script. Vous devez ajouter les méthodes save_state() et load_state() à votre script pour enregistrer et charger les états.

Pour personnaliser davantage l’emplacement et la méthode de sauvegarde des états via save_state(), utilisez la classe ProjectConfiguration. Par exemple, si automatic_checkpoint_naming est activé, chaque point de contrôle sauvegardé est stocké dans Accelerator.project_dir/checkpoints/checkpoint_{checkpoint_number}.

Tous les autres éléments avec état à stocker doivent être enregistrés via la méthode register_for_checkpointing() afin qu’ils puissent être sauvegardés et chargés. Tout objet transmis à cette méthode pour être stocké doit disposer des fonctions load_state_dict et state_dict.

Si vous avez installé torchdata>=0.8.0, vous pouvez en outre passer use_stateful_dataloader=True dans votre DataLoaderConfiguration. Cela étend les classes DataLoader d’Accelerate avec une fonction load_state_dict et une fonction state_dict, et permet à Accelerator.save_state et Accelerator.load_state de suivre également la position de lecture dans le jeu de données d’entraînement lors de la persistance du modèle.