TensorFlow モデルでの bfloat16 の使用

機械学習（ML）の研究によれば、多くの機械学習モデルでは、収束精度を低下させることなく演算の精度を低くすることが可能です。多くのモデルでは、bfloat16 を使用する場合も、32 ビットの浮動小数点数値を使用する場合と同じ収束精度で結果に到達していますが、さらに一部のモデルでは bfloat16 を使用することによって収束精度が向上しています。

このドキュメントは、bfloat16 形式を使用してメモリに活性化と勾配を格納するという点において、混合精度トレーニングに関係しています（混合精度トレーニングの詳しい背景情報については、Mixed Precision Training をご覧ください）。このドキュメントでは、Cloud TPU 上の MXU での bfloat16 の使用については説明していません。

TensorFlow を使用する ML モデルには、次のトピックが該当します。

• Google のカスタム 16 ビットの brain 浮動小数点（bfloat16）の説明。
• Cloud TPU など、ML モデルをサポートするハードウェア上の ML モデルのメモリで bfloat16 を使用することのパフォーマンス上の利点。
• TensorFlow の TPU モデルに bfloat16 を使用してメモリに活性化と勾配を格納する方法。

デフォルトでは、TensorFlow はすべての変数を 32 ビットの浮動小数点（fp32）で格納します。活性化と勾配に bfloat16 を使用すると、デバイスのステップ時間が短縮され、メモリ使用量が減少します。モデルで活性化と勾配に bfloat16 を使用する利点を判断するには、モデルの変更をご覧ください。

bfloat16 の浮動小数点形式

bfloat16 形式は、符号ビットが 1 つ、指数ビットが 8 つ、仮数ビットが 7 つ、暗黙の仮数ビットが 1 つの [1:8:7] の形式です。これに対し、標準の 16 ビット浮動小数（fp16）形式は [1:5:10] です。fp16 形式の場合、指数ビットは 5 つのみです。このような特性のため、bfloat16 の方が fp16 よりダイナミック レンジが広くなります。bfloat16 のレンジは、fp16 のダイナミック レンジの外にある可能性があるため損失のスケーリングが必要な勾配などに役立ちます。bfloat16 は、このような勾配を直接表現できます。さらに bfloat16 形式を使用すると、すべての整数 [-256, 256] を正確に表現できます。つまり、精度を失わずに int8 を bfloat16 でエンコードできます。

次の図は、3 つの浮動小数点形式を示しています。

• fp32 - IEEE の単精度浮動小数点
• fp16 - IEEE の半精度浮動小数点
• bfloat16 - 16 ビットの brain 浮動小数点

bfloat16 のダイナミック レンジは、fp16 のダイナミック レンジよりも大きくなります。

パフォーマンスとメモリ使用量の利点

Cloud TPU では、活性化や勾配などの値を bfloat16 形式で格納できます。bfloat16 を使用すると、メモリ内のデータのサイズが小さくなり、より大きいモデルを同じメモリ量に収めることができます。bfloat16 を使用すると、再実体化を削減することもできるため、ステップ時間が改善されます。

一部の演算にはメモリ帯域幅の制限があります。つまり、メモリ帯域幅によってそのような演算にかかる時間が決まります。メモリ帯域幅の制限がある演算の入出力を bfloat16 形式で格納すると、転送する必要のあるデータ量が減り、演算の速度が改善されます。

次のグラフは、社内での実験で確認した改善率を示しています。

bfloat16 を使用するようにモデルを変更する

デフォルトでは、活性化、勾配、重みはメモリに fp32 で格納されます。活性化と勾配に bfloat16 を使用し、重みは fp32 のままにして、bfloat16 と fp32 を使用した場合のモデルのパフォーマンスを比較して利点を判断できます。

1. capture_tpu_profile を使用して fp32 でモデルを実行します。

2. モデルのステップ時間と収束精度を表示するには、TensorBoard のプロファイル ビューアを使用します（詳細については、Cloud TPU ツールの使用をご覧ください）。

3. レコード パーサー内の入力パイプラインで bfloat16 に入力をキャストして、変換が input_fn の最後ではなく並列で行われるようにします。これにより、モデルのすべての活性化と勾配が bfloat16 に変換されます。

   例:

   image = tf.cast(image, tf.bfloat16)
   

4. bfloat16 のスコープでネットワークを作成し、モデルの出力を float32 にキャストします。

活性化に tf.bfloat16 を使用するようにモデルを構成した後、次の手順で bfloat16 がモデルに与える影響を確認します。

1. capture_tpu_profile を使用して、bfloat16 でモデルを実行します。
2. モデルのステップ時間と収束精度を表示するには、TensorBoard のプロファイル ビューアを使用します（詳細については、Cloud TPU ツールの使用をご覧ください）。
3. bfloat16 と fp32 のステップ時間を比較します。ステップ時間は、通常、bfloat16 の方が良好です。
4. bfloat16 と fp32 の収束精度を比較します。通常は同じですが、値が予想よりも良くなるか悪くなる可能性があります。

モデルに対して想定されるパターンの範囲がわからない場合は、複数回実行して、実行ごとの収束精度のパターンを判断することが必要になることがあります。

処理時間は速くなったが入力パイプラインがボトルネックになっていることをプロファイルが示している場合は、速度のメリットをさらに活用できるように、入力パイプラインを最適化してください。TensorFlow パイプラインのパフォーマンス向上に関する一般的なガイダンスについては、データ入力パイプラインのパフォーマンスをご覧ください。

トレーニングと推論の両方に同じ精度を使用することをおすすめします。活性化に fp32 を使用してトレーニングし、bfloat16 を使用して推論を実行することもできます（その逆も可能です）。異なる精度を選択した場合は、推論に使用された精度を使用して収束精度を検証してください。