# 混合精度トレーニング

フルプレシジョン（fp32）の学習では、すべての保存と計算が32ビットで行われます。混合精度では、計算負荷の高いフォワードパスとバックワードパスにfp16またはbf16を使用し、オプティマイザーの更新には重みのfp32コピーを保持します。これにより、計算が高速化され、重みと活性化のメモリ使用量が削減され、かつ学習の安定性が維持されます。

```text
┌─────────────────────────────────────────────────────┐
│           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 にフォールバックしてください。

> [!WARNING]
> モデルを fp32 で読み込みます。そうしないと、[autocast](https://docs.pytorch.org/docs/stable/amp.html#autocasting) は何もしなくなります。bf16 または fp16 で読み込むと、オプティマイザーが更新元の fp32 マスターコピーがなくなります。

```py
from transformers import TrainingArguments

args = TrainingArguments(..., bf16=True)
args = TrainingArguments(..., fp16=True)
```

モデルが bf16/fp16 で数値的に安定している場合、混合精度を完全にスキップして、bf16/fp16 で直接ロードおよびトレーニングを行うことができます。これにより、メモリ内の重みの fp32 コピーを回避できます。

## tf32

[tf32](https://blogs.nvidia.com/blog/tensorfloat-32-precision-format/) は Ampere GPU 上の計算モードで、matmul に 23 ビットではなく 10 ビットの mantissa を使用します。これにより高速化が得られ、特に bf16/fp16 と組み合わせる場合に効果的です。PyTorch は Ampere 以降の GPU において matmul 用の tf32 をデフォルトで有効にしていますが、[TrainingArguments](/docs/transformers/main/ja/main_classes/trainer#transformers.TrainingArguments) で明示的に設定することで、PyTorch のバージョンや環境のデフォルト設定に関係なく有効化が保証されます。

```py
from transformers import TrainingArguments

args = TrainingArguments(..., bf16=True, tf32=True)
```

## 次のステップ

- カスタム融合カーネルによる学習の高速化方法については、[Kernels](./kernels)ガイドを参照してください。
- [torch.compile](./torch_compile)のガイドを参照し、フォワードパスとバックワードパスをコンパイルしてスループットをさらに向上させる方法について学習してください。