機械学習が新参者なのでよくわからず、この辺を調べてみようと思いました。論文を読んでいたら出てきたので、ちょっと見ています。ニュースとしてはいくつか出てるので、順番に見ていきたいと思います。
Intelなどに続きArmもグーグルの「bfloat16」、完全自動運転に向けて業界団体も発足
https://xtech.nikkei.com/atcl/nxt/mag/ne/18/00007/00091/
「bfloat16は、米Google(グーグル)が独自に定義したディープラーニング向けの浮動小数点フォーマット。同社のディープラーニングアクセラレーター「TPU」で採用されている形式だ。以前は同社のみが利用していたが、米Intel(インテル)や米Facebook(フェイスブック)がサポートするなど、広がりを見せている。」
ということで、Googleの独自規格でTPUに載ってるようです。
「bfloat16では、浮動小数点の精度を示す「仮数部」に7ビットを、数値のダイナミックレンジを示す「指数部」に8ビットを割り当てる。仮数部が23ビット、指数部が8ビットのFP32(単精度浮動小数点)に比べて精度は下がるものの、指数部は8ビットで扱えるダイナミックレンジは同じにしてデータサイズを抑えたことが特徴である。ディープラーニングの演算負荷が減るほか、モデルも軽量になる。」
bfloat16 floating-point format
https://en.wikipedia.org/wiki/Bfloat16_floating-point_format
「bfloat16(brain floating point)は、コンピュータメモリーにおいて16ビットを占めるコンピュータ数値フォーマットであり、浮動小数点を用いて広い数値の動的範囲を表現します。このフォーマットは、32ビットIEEE 754単精度浮動小数点フォーマット(binary32)の短縮版(16ビット)であり、機械学習とセンサー近辺の計算を加速することを目的としています。32ビット浮動小数点数のおおよその動的範囲を維持するために、8ビットの指数部を保持していますが、binary32フォーマットの24ビット仮数に比べて8ビットの精度のみをサポートします。bfloat16数値は、単精度32ビット浮動小数点数よりもさらに、整数計算には不向きですが、それが意図された使用目的ではありません。bfloat16は、機械学習アルゴリズムの記憶要件を減らし、計算速度を上げるために使用されます。」
「bfloat16フォーマットは、Googleの人工知能研究グループであるGoogle Brainによって開発されました。Intel Xeonプロセッサー(AVX-512 BF16拡張)、Intel Data Center GPU、Intel Nervana NNP-L1000、Intel FPGAs、AMD Zen、AMD Instinct、NVIDIA GPU、Google Cloud TPUs、AWS Inferentia、AWS Trainium、ARMv8.6-A、そしてAppleのM2、それにA15チップ以降など、多くのCPU、GPU、AIプロセッサーで利用されています。多くのライブラリがbfloat16をサポートしており、CUDA、Intel oneAPI Math Kernel Library、AMD ROCm、AMD Optimizing CPU Libraries、PyTorch、TensorFlowなどがあります。これらのプラットフォームでは、bfloat16は混合精度算術においても使用されることがあり、bfloat16数値は操作され、より広いデータ型に拡張されることがあります。」
ということです。さて、本番のGoogleのTPUのbfloat16の記事を見てみます。
bfloat16 の数値形式
https://cloud.google.com/tpu/docs/bfloat16?hl=ja
「精度を落とした浮動小数点を使用して、精度を失わずに収束までの時間を短縮する方法が一般的です。TPU は、行列演算を実行するときに bfloat16 数値形式を使用します。行列乗算演算は bfloat16 値に対して実行され、累積は IEEE float32 値に対して実行されます。 bfloat16 は、1 つの符号ビット、8 つの指数ビット、7 つの仮数ビットで構成されている機械学習用のカスタム 16 ビット浮動小数点形式です。」
ということで、、、まずは符号ビット・指数ビット・仮数ビットについて確認します。
ビットで表す数字の世界~浮動小数点編~
https://www.macnica.co.jp/business/semiconductor/articles/intel/133327/
ある数字が浮動小数点計算で32bitあったときに、
-10.25 = -1.025 * 10^1
"-"が符号ビットで、プラスとマイナスで1bit
"1.025"が仮数ビットで、23bit(上記解説では32bitとなってますが、Googleの方では23bitとなっていたので、Googleの方に合わせました)
10^"1"の肩の部分が指数ビットで8bit
みたいですね。
Googleの説明に戻ります。
ウェブサイトより引用:
上記を見ると、
bfloat16は符号1bit、指数8bit、仮数7bit
fp32:単精度浮動小数点数は符号1bit、指数8bit、仮数23bit
fp16:半精度浮動小数点数は符号1bit、指数5bit、仮数10bit
ですね。明らかにfp16と異なります。指数bitはfp32と同じ8bitを維持しつつ、仮数を23bitから7bitに落としてます。
仮数部分の精度のみを下げています。それによって必要なメモリ量が半分になりますね。
「Google のハードウェア チームは、Cloud TPU 向けに bfloat16 を選択しました。これにより、float32 からの移行コストを最小限に抑えるとともに、ディープ ラーニング モデルを正確にトレーニングする能力を維持しつつ、ハードウェアの効率性を改善できます。ハードウェア乗数の物理サイズは、仮数幅の 2 乗に対応します。仮数ビットが FP16 より少ない場合、bfloat16 乗数は一般的な FP16 乗数のシリコンの約半分のサイズになり、float32 乗数の 8 分の 1 になります。」
確かに直接的な掛け算については、仮数同士の計算になるはずなので、7*7 = 49, 23*23 = 529, 10*10=100 となってかなり小さくなります。
「ニューラル ネットワークは、仮数のサイズよりも指数のサイズに敏感です。アンダーフロー、オーバーフロー、NaN で同じ動作になるように、bfloat16 の指数サイズは float32 と同じです。bfloat16 は、float32 とは異なる方法で非正規化を処理し、ゼロにフラッシュします。損失スケーリングなどの特別な処理が通常必要な float16 とは異なり、bfloat16 はディープ ニューラル ネットワークをトレーニングして実行する際の float32 の一時的な代替となります。」
ニューラルネットワーク自体が仮数よりも指数が大事みたいなので、bfloat16が精度を犠牲にせずに高速に計算できるんですね。直接的に物理的なリソースを節約できるのはとてもいいと思います。
FP8
最近はさらに浮動小数点精度を下げたFP8というのがあります。上記記事によると二つのフォーマットがあるようですが、
E5M2
・符号: 1bit
・指数: 5bit
・仮数: 2bit
コンセプトとしては、先ほどのfp32とbfloat16のように、指数bitが共通で互換性があるように、fp8はfp16と指数bitを5bitで共通化して互換性を持たせているということらしい。
E4M3
・符号: 1bit
・指数: 4bit
・仮数: 3bit
こちらは、指数が減って仮数が増えてますので、応用場面が変わりそうです。
ちなみにfp8についての論文は下記のようです。
FP8 Formats for Deep Learning
Paulius Micikevicius, Dusan Stosic, Neil Burgess, Marius Cornea, Pradeep Dubey, Richard Grisenthwaite, Sangwon Ha, Alexander Heinecke, Patrick Judd, John Kamalu, Naveen Mellempudi, Stuart Oberman, Mohammad Shoeybi, Michael Siu, Hao Wu
https://arxiv.org/abs/2209.05433
ついでに混合精度
https://www.tensorflow.org/guide/mixed_precision?hl=ja
「混合精度とは、16 ビットと 32 ビット浮動小数点型の両方を使ってモデルのトレーニングを高速化し、使用するメモリを少なくする手法です。数値の安定性を保つためにモデルの特定の部分を 32 ビットで保持することにより、モデルのステップ時間を短縮しつつ、精度などの評価指標は同様にトレーニングすることができます。このガイドでは、実験的な Keras 混合精度 API でモデルを高速化する使い方を説明します。この API を使用すると、最新の GPU で 3 倍以上、TPU で 60% 以上パフォーマンスを向上させることができます。」
基本的にはこうしたデータのタイプに対応しているハードウェアを使う必要がありますが、ディープラーニングにあった形のデータを使うことでいろいろ実用的な計算がより高速にできます。互換性を考えたりとかなり工夫されていますので、その辺りを意識してソフトを組んだりハードを選ぶと楽しそうです。以上です。