bfloat16 を使用してモデルのパフォーマンスを向上させる

デフォルトでは、TPU は bfloat16 値で行列乗算演算を実行し、IEEE float32 値で累積を実行します。精度を落とした浮動小数点数を使用すると、精度を失わずに収束までの時間を短縮できます。

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

bfloat16 を明示的に使用する

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

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

開始するには、Cloud TPU リファレンスモデルのいずれかについて経験を積むことをおすすめします。その後、プロファイリングツールのガイドとトラブルシューティングガイドに記載された詳細な技術情報を使用して、機械学習モデルを独自に作成し、最適化することができます。

形式変換の詳細

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

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

Cloud TPU でトレーニングされたモデルから取得したチェックポイントは、大規模な手動変換なしで、他のハードウェアプラットフォームにデプロイできます（CPU や GPU での推論や微調整など）。