Améliorer les performances de votre modèle avec bfloat16

Par défaut, les TPU effectuent des opérations de multiplication matricielle avec des valeurs bfloat16 et des accumulations avec des valeurs IEEE float32. L'utilisation de nombres à virgule flottante de précision réduite réduit le délai de sans perdre en précision.

La plage dynamique de bfloat16 et float32 est équivalente. Toutefois, bfloat16 utilise la moitié de l'espace mémoire. Pour en savoir plus sur les performances de bfloat16, consultez A Study of BFLOAT16 for Deep Learning Training (Étude sur BFLOAT16 pour l'entraînement en deep learning).

Utiliser bfloat16 explicitement

Bien que la conversion automatique de format dans les TPU vous permette d'éviter de penser à la précision numérique, vous pouvez améliorer les performances en convertissant explicitement des valeurs en bfloat16. Il existe deux raisons de caster explicitement des valeurs en bfloat16 :

1. Le stockage de valeurs au format bfloat16 permet d'économiser de la mémoire sur puce, permettant ainsi aux Cloud TPU d'entraîner des modèles plus volumineux ou d'utiliser des tailles de lot plus importantes.

2. Certaines opérations sont associées à la bande passante mémoire, ce qui signifie que le temps nécessaire au chargement des données à partir de la mémoire peut ralentir le temps global d'exécution du calcul. Le stockage des opérandes et des sorties de ces opérations au format bfloat16 réduit la quantité de données à transférer, améliorant ainsi la vitesse globale.

Pour commencer, nous vous recommandons d'acquérir une expérience avec l'un des modèles de référence Cloud TPU. Après le guide des outils de profilage. guide de dépannage fournissent Des informations techniques détaillées pour vous aider à créer et à optimiser le machine learning des modèles par vous-même.

Informations sur la conversion de format

La conversion de format de float32 vers bfloat16 est automatiquement insérée par le compilateur XLA. Sur TPU, le schéma d'arrondi dans la conversion est arrondi au nombre pair le plus proche et déborder sur inf. De plus, bfloat16 sur Cloud TPU n'est pas compatible avec les valeurs sous-normales. Par conséquent, toutes les valeurs sous-normales sont mises à zéro lors de la conversion. Les valeurs spéciales, telles que NaN et inf, sont conservées dans la conversion.

La conversion du format de bfloat16 en float32 est également insérée automatiquement par le compilateur XLA. Étant donné que float32 peut représenter toutes les valeurs exactes dans bfloat16, la conversion comble 16 zéros dans les bits de mantisse. Les valeurs spéciales sont conservées dans la conversion.

Les points de contrôle obtenus à partir d'un modèle entraîné sur Cloud TPU peuvent être déployés sur d'autres plates-formes matérielles (par exemple, des inférences ou des réglages sur des processeurs ou des GPU) sans conversions manuelles approfondies.