Entraînement sur pods TPU

Présentation

Les TPU ont été conçus pour être adaptés aux pods TPU. Un pod TPU est un ensemble d'appareils TPU connectés par des interfaces réseau haut débit dédiées. Un pod TPU vous permet de répartir la charge de traitement sur plusieurs TPU. Chaque carte TPU est connectée à une machine hôte hautes performances basée sur le processeur pour des opérations, telles que le chargement et le prétraitement des données. Pour tirer pleinement parti d'un grand nombre de TPU, vous devez régler plusieurs paramètres de tâche d'entraînement.

La configuration de l'entraînement avec des pods TPU est différente pour chaque framework. Utilisez les liens suivants pour consulter des informations détaillées concernant l'entraînement sur les pods avec chaque framework:

Les sections suivantes décrivent certains problèmes courants, les modifications à apporter à vos modèles et les bonnes pratiques à suivre pour réduire ou éviter les défaillances des pods.

Adapter la taille de lot et les pas d'entraînement

Pour effectuer un scaling linéaire sur des TPU de plus grande taille, ne modifiez pas la taille de lot par cœur.

Par exemple, si vous utilisez une taille de lot de 1 024 sur un type v2-8, utilisez une taille de lot de 4 096 (4 * 1 024) sur un type v2-32. Cela permet d'utiliser pleinement le matériel TPU. Vous pouvez utiliser des tailles de lot plus petites, mais dans ce cas, la mise à l'échelle de l'entraînement ne sera pas linéaire.

Certains modèles incluent un indicateur train_steps dans lequel une étape correspond au traitement d'un seul lot de données. Lorsque vous augmentez la taille du lot, réduisez le nombre de pas d'entraînement afin que le nombre total d'exemples d'entraînement reste le même.

Par exemple, si la taille de lot est de 1 000 pour 100 pas, 100 000 exemples sont traités pendant l'entraînement. Si vous avez maintenant quatre nœuds de calcul et une taille de lot effective de 4 000, vous devrez ajuster le nombre de pas sur 25 pour traiter ces 100 000 exemples. Si le modèle utilise l'option epochs, vous n'avez pas besoin d'adapter le nombre de pas.

Les tailles de lot plus grandes pouvant modifier le comportement de convergence du modèle, vous devrez peut-être ajuster certains hyperparamètres, tels que le taux d'apprentissage.

Utiliser des buckets Google Cloud Storage régionaux dans la même région que le pod TPU

En règle générale, pour l'entraînement TPU, il est recommandé de toujours utiliser les ressources dans la même région. La région des ressources est particulièrement importante lorsque vous utilisez des pods TPU, car les taux de transfert de données sont plus élevés lorsque votre bucket Google Cloud Storage et votre TPU se trouvent dans la même région.

Vérifiez que vous utilisez un bucket Google Cloud Storage régional dans la même région que le TPU pour les ensembles de données d'entraînement et les points de contrôle.

Bonnes pratiques de workflow pour le développement sur des pods TPU

Lors du développement d'une nouvelle charge de travail TPU, il est souvent préférable de commencer à développer des TPU de taille réduite et d'effectuer une itération progressive vers de plus grandes tailles. Commencez par utiliser une petite version de TPU (par exemple, v2-8 ou v3-8).

  • Tester le fonctionnement de votre charge de travail
  • Tester et valider les performances à l'aide d'outils dédiés

Une fois que votre charge de travail est fonctionnelle et atteint vos objectifs de performances, passez à un type de TPU plus grand, tel que v2-32 ou v3-32. Augmentez progressivement et de manière itérative la taille de TPU tout en validant l'évolutivité (fonctionnalités et performances) jusqu'à atteindre la taille de TPU souhaitée.