Utiliser le format bfloat16 avec les modèles TensorFlow

La recherche en Machine learning (ML) montre que de nombreux modèles de machine learning peuvent tolérer une arithmétique de précision inférieure sans dégradation de la précision convergée. Avec bfloat16, de nombreux modèles obtiennent des résultats identiques à ceux obtenus à l'aide des valeurs numériques à virgule flottante de 32 bits, et certains modèles présentent même une précision convergente améliorée avec bfloat16.

Ce document s'intéresse à l'entraînement de précision mixte au sens du stockage en mémoire d'activations et de gradients au format bfloat16. (Pour obtenir plus d'informations sur l'entraînement de précision mixte, consultez la section Entraînement de précision mixte). Ce document n'aborde pas l'utilisation de bfloat16 dans le MXU sur Cloud TPU.

Les rubriques suivantes s'appliquent aux modèles ML utilisant TensorFlow :

  • Description de bfloat16, le format à virgule flottante brain de 16 bits de Google.
  • Avantages de l'utilisation de bfloat16 en mémoire pour les modèles ML sur le matériel compatible, tel que Cloud TPU.
  • Comment stocker les activations et les gradients en mémoire à l'aide de bfloat16 pour un modèle TPU dans TensorFlow.

Par défaut, TensorFlow stocke toutes les variables au format 32 bits à virgule flottante (fp32). L'utilisation de bfloat16 pour les activations et les gradients accélère le temps de processeur de l'appareil et réduit l'utilisation de la mémoire. Consultez la section Modifier votre modèle pour déterminer les avantages de l'utilisation de bfloat16 pour les activations et les gradients dans votre modèle.

Le format à virgule flottante bfloat16

Le format bfloat16 est [1:8:7], qui comprend un bit de signe, huit bits d'exposant, sept bits de mantisse et un bit de mantisse implicite. En comparaison, le format standard à virgule flottante 16 bits (fp16) est [1:5:10]. Notez que le format fp16 ne comprend que cinq bits d'exposant. En raison de ces caractéristiques, bfloat16 dispose d'une plage dynamique supérieure à celle de fp16. La plage bfloat16 est utile pour des éléments tels que les dégradés qui peuvent se trouver en dehors de la plage dynamique de fp16 et nécessitent donc une mise à l'échelle des pertes. bfloat16 peut représenter ces dégradés directement. En outre, vous pouvez utiliser le format bfloat16 pour représenter avec précision tous les nombres entiers [-256, 256], ce qui signifie que vous pouvez encoder des données int8 au format bfloat16 sans perte de précision.

La figure suivante affiche trois formats de nombre à virgule flottante

  • fp32 - virgule flottante à simple précision IEEE
  • fp16 - virgule flottante demi-précision IEEE
  • bfloat16 - virgule flottante brain 16 bits

image

La plage dynamique de bfloat16 est supérieure à celle de fp16.

Avantages de performance et d'utilisation de la mémoire

Cloud TPU est compatible avec le stockage de valeurs telles que les activations et les gradients au format bfloat16. L'utilisation de bfloat16 réduit la taille des données en mémoire et permet aux modèles plus grands de tenir dans la même quantité de mémoire. L'utilisation de bfloat16 peut également réduire la rematérialisation, ce qui optimise le temps de processeur.

Certaines opérations sont associées à la bande passante mémoire, ce qui signifie que cette dernière détermine le temps passé sur ces opérations. Le stockage des entrées et des sorties des opérations associées à la bande passante mémoire au format bfloat16 réduit la quantité de données à transférer, améliorant ainsi la vitesse des opérations.

Le graphique suivant affiche les améliorations observées lors de nos expériences internes.

image

Modifier votre modèle pour utiliser le format bfloat16

Par défaut, les activations, les gradients et les pondérations sont stockés en mémoire au format fp32. Vous pouvez utiliser bfloat16 pour les activations et les gradients en laissant des pondérations au format fp32, puis observer les performances de votre modèle avec bfloat16 et fp32 pour déterminer les avantages comparatifs.

  1. Exécutez le modèle au format fp32 en utilisant capture_tpu_profile.

  2. Pour afficher le temps de processeur et la précision convergée du modèle, utilisez la visionneuse de profil dans TensorBoard. (Consultez la section Utilisation des outils Cloud TPU pour plus de détails).

  3. Saisissez l'entrée au format bfloat16 dans votre pipeline d'entrées dans l'analyseur d'enregistrement afin que la conversion puisse être effectuée en parallèle plutôt qu'à la fin de input_fn. Cette opération convertit toutes les activations et tous les gradients du modèle en bfloat16.

    Par exemple :

    image = tf.cast(image, tf.bfloat16)
    
  4. Créez votre réseau sous le champ d'application bfloat16, puis saisissez les sorties du modèle au format float32.

Une fois que vous avez configuré votre modèle pour utiliser tf.bfloat16 pour les activations, vérifiez les éléments suivants pour voir l'impact de bfloat16 sur votre modèle :

  1. Exécutez le modèle avec bfloat16 à l'aide de capture_tpu_profile.
  2. Pour afficher le temps de processeur et la précision convergée du modèle, utilisez la visionneuse de profil dans TensorBoard. (Consultez la section Utilisation des outils Cloud TPU pour plus de détails).
  3. Comparez le temps de processeur pour le format bfloat16 et le format fp32. Le temps de processeur connaît en général une amélioration avec le format bfloat16.
  4. Comparez la précision convergée pour le format bfloat16 par rapport au format fp32. Elles sont généralement identiques, mais les valeurs peuvent être meilleures ou pires que prévu.

Si vous ne connaissez pas la plage de variation à prévoir pour votre modèle, vous allez peut-être devoir répéter les exécutions pour déterminer la variation de la précision convergée.

Si votre profil indique que le temps de traitement est plus court, mais que le pipeline d'entrée est devenu un goulot d'étranglement, optimisez le pipeline pour améliorer la vitesse. Consultez la section Performances du pipeline d'entrée de données pour obtenir des instructions générales sur l'amélioration des performances du pipeline TensorFlow.

Un exemple de bonne pratique pour l'entraînement et l'inférence consiste à utiliser la même précision pour les deux processus. Il est possible de s'entraîner à l'aide du format fp32 pour les activations, puis d'exécuter l'inférence avec le format bfloat16 (ou inversement). Si vous choisissez une précision incompatible, vérifiez la précision convergée à l'aide de la précision utilisée pour l'inférence.