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 à précision réduite réduit le temps de convergence sans perdre en précision.

La plage dynamique de bfloat16 et float32 est équivalente. Toutefois, bfloat16 occupe 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. Ensuite, le guide des outils de profilage et le guide de dépannage fournissent des informations techniques détaillées qui vous aideront à créer et à optimiser des modèles de machine learning 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 un TPU, le schéma d'arrondi de la conversion est arrondi à l'entier pair le plus proche et le débordement est défini 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 définies sur zéro lors de la conversion. Les valeurs spéciales, telles que NaN et inf, sont conservées lors de la conversion.

La conversion de format de bfloat16 vers 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 ajoute 16 zéros dans les bits de la mantisse. Les valeurs spéciales sont conservées lors de 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.