Mixed Precision

混合精度は、ニューラルネットワークの計算の大部分を、FP16（半精度）またはBF16のいずれかの16ビット浮動小数点形式で実行しながら、モデルの重みの32ビット（FP32）マスターコピーを保持し、数値精度が特に重要な演算はFP32で実行する学習手法です。16ビット値はFP32に比べてメモリ使用量と帯域幅が半分で済み、最新のTensor Core搭載GPUでは16ビット値をはるかに高速に乗算できるため、混合精度では学習時間を短縮し、メモリ内により大きなバッチサイズを収めつつ、FP32と同等の精度を維持できます。2018年にMicikeviciusらが説明した、FP32のマスター重み、FP16演算、損失スケーリングを組み合わせる手法は、現在、ディープラーニングフレームワークでは自動混合精度（AMP）として自動化されています。

混合精度の仕組み#

16ビット学習での誤差逆伝播と重みの更新を、FP32学習と同等の精度にするには、次の3つの要素が重要です。

1. FP32マスター重み：オプティマイザーは、重みのFP32コピーに各更新を適用します。FP16では有効桁数が約3桁しかないため、小さな更新値をFP16の重みに直接加えると、多くの場合、丸めによって消えてしまいます。
2. 演算ごとの精度設定：autocastコンテキストでは、Tensor Coreによる高速化が最も大きい行列乗算と畳み込みを16ビットで実行し、リダクション、ソフトマックス、損失はFP32で保持します。
3. 損失スケーリング：FP16では約6 × 10⁻⁸未満の値を表現できないため、多くの小さな勾配がアンダーフローしてゼロになります。誤差逆伝播の前に損失へスケール係数を掛けることで、すべての勾配をFP16の範囲内に移動できます。重みの更新前には、勾配を同じ係数で割ります。

PyTorch AMPでは、torch.autocastが演算ごとの型変換を行い、GradScalerが損失スケーリングを行います。

import torch
from torch import nn

# Use FP16 on an NVIDIA GPU and BF16 on CPU
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if device == "cuda" else torch.bfloat16

model = nn.Linear(512, 10).to(device)  # master weights stay in FP32
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
scaler = torch.amp.GradScaler(device, enabled=dtype == torch.float16)  # loss scaling for FP16 only

x = torch.randn(64, 512, device=device)
with torch.autocast(device_type=device, dtype=dtype):
    out = model(x)  # the matrix multiply runs in 16-bit
    loss = out.float().pow(2).mean()

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
print(f"Output: {out.dtype}, weights: {model.weight.dtype}, gradients: {model.weight.grad.dtype}")

CPUでは、このコードはレイヤーの出力にtorch.bfloat16、重みと勾配の両方にtorch.float32を出力します。NVIDIA GPUでは、出力はtorch.float16です。1回の学習ステップで、16ビット演算とFP32のマスター重みが共存します。

PyTorchのGradScalerは、スケールを動的に調整します。65,536から開始し、スケーリングされた勾配にinfまたはNaNの値が含まれるたびにスケールを半分にしてステップをスキップします。オーバーフローせずに2,000ステップ連続で進むとスケールを2倍にするため、勾配がオーバーフローする直前の値に係数が落ち着きます。

FP16、BF16、さらに低い精度#

FP16では指数部に5ビット、仮数部に10ビットを割り当てるため精度は高い一方、65,504を超えるとオーバーフローします。BF16はFP32と同じ8ビットの指数部を保持し、約3.4 × 10³⁸まで同じ範囲をカバーしますが、仮数部は7ビットしかないため、丸め精度が粗くなる代わりに損失スケーリングが不要です。BF16はAmpere世代以降のNVIDIA GPU、GoogleのTPU、および多くの最新CPUで動作し、大規模言語モデルの一般的な学習形式です。同じ考え方は、NVIDIAのTransformer Engineを通じて、HopperおよびBlackwell GPUのFP8にも適用されます。また、BlackwellではNVFP4などのFP4形式にも広がっています。

混合精度の実践#

AMPの主な実用上のリスクは、16ビットでNaN損失を生じさせるモデルやGPUです。Ultralytics YOLO26の学習ではAMPがデフォルトで有効になっており、amp=Trueを使用します。まずCUDA GPU上でYOLO26nモデルのFP32予測とAMP予測を比較し、結果が一致しない場合、またはGTX 16シリーズのようにAMPでNaN損失が発生することが知られているシリーズのGPUの場合は、FP32に切り替えます。amp="bf16"を設定すると、損失スケーリングを無効にしてBF16 autocastに切り替わります。

混合精度は、学習がメモリや時間によって制約される場面で広く利用されています。自動運転車の開発チームは、数百万枚の高解像度フレームで物体検出モデルを学習します。活性化メモリを半減できれば、より大きな画像やバッチを扱えます。医用画像解析では、3D CTやMRIのボリュームがFP32ではGPUメモリを超えることが多く、16ビットの活性化を使うことで、ボリューム全体を対象としたセグメンテーションモデルの学習が可能になります。

混合精度と関連概念の比較#

• モデルの量子化：量子化では、学習済みモデルをINT8などの低精度整数に変換して推論を高速化します。一方、混合精度では浮動小数点の重みを維持しながら学習を高速化します。
• 量子化を考慮した学習：QATは学習中に整数の丸めをシミュレーションし、INT8でデプロイした後もモデルの精度を維持できるようにします。
• 半精度：半精度はFP16の数値形式そのものを指し、混合精度はFP16とFP32を組み合わせる学習手法を指します。FP32のマスター重みや損失スケーリングを使わずにFP16だけで学習すると、小さな更新値は丸めによって失われ、小さな勾配はアンダーフローし、収束しないこともよくあります。