bfloat16 の数値形式

精度を落とした浮動小数点を使用して、精度を失わずに収束までの時間を短縮する方法が一般的です。TPU は、行列演算を実行するときに bfloat16 数値形式を使用します。行列乗算演算は bfloat16 値に対して実行され、累積は IEEE float32 値に対して実行されます。

bfloat16 は、1 つの符号ビット、8 つの指数ビット、7 つの仮数ビットで構成されている機械学習用のカスタム 16 ビット浮動小数点形式です。　　次の図は、float32: IEEE の単精度、float16: IEEE の半精度、bfloat16 の 3 つの浮動小数点形式の内部を示しています。

bfloat16 と float32 のダイナミックレンジは同等です。ただし、bfloat16 はメモリ容量の半分を消費します。bfloat16 のパフォーマンスの詳細については、ディープラーニングトレーニングのための BFLOAT16 の研究をご覧ください。

bfloat16 の選択

Google のハードウェアチームは、Cloud TPU 向けに bfloat16 を選択しました。これにより、float32 からの移行コストを最小限に抑えるとともに、ディープラーニングモデルを正確にトレーニングする能力を維持しつつ、ハードウェアの効率性を改善できます。ハードウェア乗数の物理サイズは、仮数幅の 2 乗に対応します。仮数ビットが FP16 より少ない場合、bfloat16 乗数は一般的な FP16 乗数のシリコンの約半分のサイズになり、float32 乗数の 8 分の 1 になります。

ニューラルネットワークは、仮数のサイズよりも指数のサイズに敏感です。アンダーフロー、オーバーフロー、NaN で同じ動作になるように、bfloat16 の指数サイズは float32 と同じです。bfloat16 は、float32 とは異なる方法で非正規化を処理し、ゼロにフラッシュします。損失スケーリングなどの特別な処理が通常必要な float16 とは異なり、bfloat16 はディープニューラルネットワークをトレーニングして実行する際の float32 の一時的な代替となります。

混合精度トレーニング

ディープニューラルネットワーク内のほとんどの計算では、低精度の値を使用して同じ精度でタスクを達成できます。一部のモデルでは、低い精度の値で高い精度に到達することもできます。

Cloud TPU をプログラミングする場合、XLA コンパイラは float32 と bfloat16 の間で値を自動的に変換します。

フォーマット変換の詳細

float32 から bfloat16 への形式変換は、XLA コンパイラによって自動的に挿入されます。TPU では、変換の丸めスキームは最も近い偶数に丸められ、inf にオーバーフローします。また、Cloud TPU の bfloat16 は非正規化をサポートしていないため、変換中にすべての非正規化がゼロにフラッシュされます。 NaN や inf などの特殊な値は変換で保持されます。

bfloat16 から float32 への形式変換も、XLA コンパイラによって自動的に挿入されます。float32 は bfloat16 のすべての正確な値を表すことができるため、変換では、仮数ビットに 16 のゼロがパディングされます。変換では特殊な値が保持されます。

モデルのポータビリティ

TPU ハードウェアはこれらの値を bfloat16 に自動的にキャストできるため、モデル内のパラメータの値と有効化は 32 ビット形式で保存できます。Cloud TPU でトレーニングされたモデルから取得したチェックポイントは、大規模な手動変換なしで、他のハードウェアプラットフォームにデプロイできます（CPU や GPU での推論や微調整など）。

bfloat16 によるパフォーマンスの改善

TPU での自動形式変換により、数値の精度について考慮する必要がなくなりますが、値を明示的に bfloat16 にキャストすることで、さらにパフォーマンスを向上させることができます。値が明示的に bfloat16 にキャストされる理由は 2 つあります。

値を bfloat16 形式で格納するとオンチップメモリを節約できるため、Cloud TPU でより大きなモデルのトレーニングや、より大きなバッチサイズの使用が可能になります。
一部の演算にはメモリ帯域幅の制限があります。つまり、メモリからデータを読み込むのにかかる時間が、計算の実行に費やす全体的な速度を低下させる可能性があります。これらの演算のオペランドと出力を bfloat16 形式で格納すると、転送する必要があるデータ量が減るため、全体的な速度が向上します。

開始するには、Cloud TPU 用に最適化された bfloat16 対応のリファレンスモデルのいずれかでハンズオン経験を行うことをおすすめします。Google のパフォーマンスガイド、プロファイリングツールガイド、トラブルシューティングガイドに、マシンの作成と最適化に役立つ詳細な技術情報が記載されています。ご自身で機械学習モデルを作成して最適化してください。