Entraînement sur TPU

Un TPU (Tensor Processing Unit) est un type de matériel conçu spécifiquement pour entraîner des modèles de manière efficace. Accelerate prend en charge l'entraînement sur TPU, mais il y a quelques points à garder à l'esprit, notamment la compilation de graphes. Ce tutoriel aborde brièvement la compilation ; pour plus de détails, consultez le guide Training on TPUs with Accelerate.

Compilation

Un TPU crée un graphe de toutes les opérations de l'étape d'entraînement, telles que le passage avant, le passage arrière et l'étape de l'optimiseur. C'est pourquoi la première étape d'entraînement prend toujours un certain temps, car la construction et la compilation de ce graphe sont chronophages. Mais une fois la compilation terminée, elle est mise en cache et toutes les étapes suivantes sont beaucoup plus rapides.

La clé est d'éviter de recompiler votre code, sinon l'entraînement est extrêmement lent. Cela signifie que toutes vos opérations doivent être exactement les mêmes :

• tous les tenseurs de vos lots doivent avoir la même longueur (par exemple, aucun bourrage dynamique pour les tâches de NLP)
• votre code doit être statique (par exemple, pas de couches avec des boucles for dont la longueur varie selon l'entrée, comme un LSTM)

Liaison des poids

Une conception courante des modèles de langage consiste à lier les poids des couches d’incorporation et de softmax. Toutefois, le transfert du modèle sur un TPU (que ce soit vous-même ou en le passant à la méthode prepare()) rompt le lien entre les poids, et vous devrez alors relier les poids à nouveau.

Pour ajouter un comportement spécial (comme le partage de poids) dans votre script pour les TPU, définissez d'abord distributed_type sur DistributedType.TPU. Vous pouvez ensuite utiliser la méthode tie_weights pour lier les poids.

if accelerator.distributed_type == DistributedType.TPU:
    model.tie_weights()