Transformers documentation
混合精度トレーニング
混合精度トレーニング
フルプレシジョン（fp32）の学習では、すべての保存と計算が32ビットで行われます。混合精度では、計算負荷の高いフォワードパスとバックワードパスにfp16またはbf16を使用し、オプティマイザーの更新には重みのfp32コピーを保持します。これにより、計算が高速化され、重みと活性化のメモリ使用量が削減され、かつ学習の安定性が維持されます。
┌─────────────────────────────────────────────────────┐ │ MIXED PRECISION TRAINING LOOP │ │ │ │ fp32 master weights ──cast──▶ fp16/bf16 │ │ ▲ │ │ │ │ FORWARD (autocast) │ │ │ matmuls in fp16/bf16 │ │ │ reductions stay fp32 │ │ │ │ loss │ │ │ LOSS SCALE ×S ──fp16 │ │ │ │ │ │ │ BACKWARD │ │ │ grads in fp16/bf16 │ │ │ │ │ │ │ UNSCALE ÷S ──fp16 │ │ │ check inf/nan │ │ │ cast grads → fp32 │ │ └────────────────────────── optimizer.step │ └─────────────────────────────────────────────────────┘
bf16() または fp16() を True に設定すると、混合精度トレーニングが有効になります。どちらも 16 ビット型ですが、bf16 は fp32 と同じ指数範囲を持つため、オーバーフローがほぼ発生しません。Ampere 世代以降の GPU（A100、H100 など）では bf16 を使用し、V100 や T4 などの旧世代ハードウェアでは fp16 にフォールバックしてください。
モデルを fp32 で読み込みます。そうしないと、autocast は何もしなくなります。bf16 または fp16 で読み込むと、オプティマイザーが更新元の fp32 マスターコピーがなくなります。
from transformers import TrainingArguments
args = TrainingArguments(..., bf16=True)
args = TrainingArguments(..., fp16=True)モデルが bf16/fp16 で数値的に安定している場合、混合精度を完全にスキップして、bf16/fp16 で直接ロードおよびトレーニングを行うことができます。これにより、メモリ内の重みの fp32 コピーを回避できます。
tf32
tf32 は Ampere GPU 上の計算モードで、matmul に 23 ビットではなく 10 ビットの mantissa を使用します。これにより高速化が得られ、特に bf16/fp16 と組み合わせる場合に効果的です。PyTorch は Ampere 以降の GPU において matmul 用の tf32 をデフォルトで有効にしていますが、TrainingArguments で明示的に設定することで、PyTorch のバージョンや環境のデフォルト設定に関係なく有効化が保証されます。
from transformers import TrainingArguments
args = TrainingArguments(..., bf16=True, tf32=True)次のステップ
- カスタム融合カーネルによる学習の高速化方法については、Kernelsガイドを参照してください。
- torch.compileのガイドを参照し、フォワードパスとバックワードパスをコンパイルしてスループットをさらに向上させる方法について学習してください。