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
の間で値を自動的に変換します。
モデルのポータビリティ
TPU ハードウェアはこれらの値を bfloat16
に自動的にキャストできるため、モデル内のパラメータの値と有効化は 32 ビット形式で保存できます。Cloud TPU でトレーニングされたモデルから取得したチェックポイントは、大規模な手動変換なしで、他のハードウェア プラットフォームにデプロイできます(CPU や GPU での推論や微調整など)。
bfloat16 によるパフォーマンスの改善
TPU での自動形式変換により、数値の精度について考慮する必要がなくなりますが、値を明示的に bfloat16
にキャストすることで、さらにパフォーマンスを向上させることができます。値が明示的に bfloat16
にキャストされる理由は 2 つあります。
値を
bfloat16
形式で格納するとオンチップ メモリを節約できるため、Cloud TPU でより大きなモデルのトレーニングや、より大きなバッチサイズの使用が可能になります。一部の演算にはメモリ帯域幅の制限があります。つまり、メモリからデータを読み込むのにかかる時間が、計算の実行に費やす全体的な速度を低下させる可能性があります。これらの演算のオペランドと出力を
bfloat16
形式で格納すると、転送する必要があるデータ量が減るため、全体的な速度が向上します。
開始するには、Cloud TPU 用に最適化された bfloat16
対応のリファレンス モデルのいずれかでハンズオン経験を行うことをおすすめします。Google のパフォーマンス ガイド、プロファイリング ツールガイド、トラブルシューティング ガイドに、マシンの作成と最適化に役立つ詳細な技術情報が記載されています。ご自身で機械学習モデルを作成して最適化してください。