混合精度学習 — fp16/bf16/fp8で壊れずに速くする

同じモデルを半分のビット数で学習させると、速くなる代わりに勾配が0やinfに化けます。壊れ方の正体（アンダーフローとオーバーフロー）から出発し、loss scalingが何をしている一手なのか、bf16はなぜそれを不要にしたのか、fp8で何が戻ってきたのかを順に。最後はPyTorch AMPの設定と、クリップ順序を間違えて静かに事故る典型例まで。

Mixed Precision Training

粗いものさしで測ると、何が起きるか

工具箱に巻尺が2本あるとします。1本は1mm刻みで、測れるのは3mまで。もう1本は1cm刻みで、30mまで届く。どちらも目盛りの数はほぼ同じです。細かく測るか、遠くまで測るかは交換条件で、目盛りの総数（＝ビット数）が決まっている以上、片方を取れば片方を失います。

コンピュータの数値もこれと同じです。32ビットを16ビットに減らせば、メモリは半分、GPUの行列積は専用回路（Tensor Core）に乗って速くなる。その代わり、目盛りが粗くなるか、測れる範囲が狭くなるか、あるいは両方が起きます。

問題は、学習中に流れる数の大きさが恐ろしく幅広いことです。損失は1前後、重みは0.01前後、勾配は0.0001前後、ときに0.00000001。この幅がものさしからはみ出した瞬間、値は静かに 0 か inf になります。混合精度学習とは、この事故を起こさずに粗いものさしを使い切るための工夫の総称です。

「混合」とは、3つの役割を別の型に割り振ること

名前が「混合」なのは、モデル全体を低精度にするわけではないからです。学習中の数値には少なくとも3つの役割があり、必要な精度がそれぞれ違います。

1. 計算する（行列積・畳み込み）— 入力は低精度でよい。ここが速度のほぼ全部
2. 足し込む（累積・総和）— 何千個も足すので、途中で桁が落ちると効いてくる
3. 保管する（重みそのもの）— 何万ステップも少しずつ更新されるので、細かい差を捨ててはいけない

混合精度は、この3つを別々の型に割り振ります。計算はfp16やbf16、累積はfp32、保管もfp32。GPUの行列積回路は、もともと低精度で掛けてfp32で足す設計になっているので、1と2は勝手にそうなります。書き手が意識して守るのは3つ目、いわゆるマスター重みです。

2017年の Mixed Precision Training はまさにこの3点セット（fp32のマスター重み・loss scaling・fp32での累積）を提案した論文で、画像分類から機械翻訳、音声認識まで、fp32と同等の精度を保てることを示しました。以降の学習フレームワークが既定で持っている仕組みは、ほぼこの論文の延長線上にあります。

学習が壊れる2つの落ち方

fp16が表せるのは、ざっくり 0.00006 から 65504 までの範囲です（負の側は対称）。窓の外に出た値は2通りの壊れ方をします。

上にはみ出す（オーバーフロー） と inf になります。inf は次の演算で inf - inf などを作って nan に化け、nan は掛けても足しても nan のままなので、1個生まれると数ステップでモデル全体に伝染します。損失が突然 nan になって二度と戻らないのはこれです。派手に壊れるぶん、気づけます。

下にはみ出す（アンダーフロー） はもっと厄介で、値は静かに 0 になります。エラーも警告も出ません。勾配が0になった重みは更新されず、そのまま学習が進み、「なぜか精度が伸びない」という形でしか表に出ません。原論文が示した勾配のヒストグラムでは、勾配の大きさのかなりの部分がfp16の下限より小さい側に落ちていました。放っておけば消えるのが普通、というのが出発点です。

厄介なのは、これが学習率の問題と見分けづらいことです。更新が小さすぎれば進まないし、大きすぎれば飛ぶ——という当たり前の挙動を、下の図で先に体に入れておいてください。混合精度で起きるのは、学習率をいじっていないのに、数の表現だけが原因で更新が0やinfになるという事態です。

指数部と仮数部 — fp16とbf16は同じ16ビットの別解

なぜ0や inf になるのかは、ビットの割り振りを見ると一発で分かります。浮動小数点数は3つの区画でできています。

$s$ は符号（0なら正、1なら負）、$e$ は指数部に入る整数、$b$ はその下駄（バイアス）で型ごとに決まった定数、$f$ は仮数部が表す0以上1未満の小数です。つまり「2の何乗か」で大きさのケタを決め、「1.○○○」の部分で細かさを決めている、という形です。指数部のビット数が測れる範囲を、仮数部のビット数が目盛りの細かさを決めます。

同じ16ビットでも、fp16とbf16は真逆の削り方をしています。fp16は指数部を8→5に削って範囲を捨て、仮数部を10ビット残して細かさを取った。bf16は指数部を8のまま据え置いて範囲をfp32と揃え、仮数部を7ビットまで削って細かさを捨てた。ビット配置の詳しい読み方は数の表し方 — FP32からFP8・INT4までで扱っています。