Entraîner un modèle sur des tranches TPU

Les TPU sont conçus pour être évolutifs et s'adapter 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 les CPU 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 obtenir des informations détaillées concernant l'entraînement sur des pods avec chaque framework :

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

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

Pour effectuer un scaling linéaire sur des types de TPU plus grands, conservez la même taille de lot par cœur.

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

Certains modèles incluent un flag train_steps dans lequel un pas 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 le flag 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 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 Cloud Storage et votre TPU se trouvent dans la même région.

Vérifiez que vous utilisez un bucket 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 sur des TPU de taille réduite et d'effectuer une itération progressive vers des TPU de plus grande taille. Commencez par utiliser une version de TPU de petite taille (par exemple, v6e-8).

  • Testez le comportement attendu de votre charge de travail
  • Testez et validez les performances à l'aide d'outils dédiés

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