混合精度演算(Mixed Precision Training)とは、ディープラーニングの学習において、FP32(32ビット)とFP16またはBFloat16(16ビット)など、異なる精度の数値フォーマットを組み合わせて使用する手法です。学習速度を大幅に向上させながら、モデルの精度をほぼ維持できるため、現在の大規模モデル学習では事実上の標準となっています。
混合精度演算の仕組み
混合精度学習では、順伝播と逆伝播の行列演算をFP16/BFloat16で実行し、パラメータのマスターコピーはFP32で保持します。FP16の演算はTensor Coreなどの専用ハードウェアで高速に処理され、FP32のマスターコピーにより累積的な丸め誤差を防ぎます。損失スケーリング(Loss Scaling)によりFP16のアンダーフローを回避します。
性能上のメリット
混合精度演算により、FP32のみの学習と比較して2〜3倍の高速化が実現できます。メモリ使用量も大幅に削減されるため、より大きなバッチサイズやモデルサイズでの学習が可能になります。NVIDIA GPUのTensor Coreは混合精度演算に最適化されており、最大限の性能を引き出せます。
実装と対応フレームワーク
PyTorchのtorch.cuda.amp、TensorFlowのtf.keras.mixed_precision、NVIDIA Apexなどにより簡単に混合精度学習を導入できます。BFloat16はGPUやTPUでの学習で広く使われており、損失スケーリングなしでも安定した学習が可能なため、実装がさらにシンプルです。