B-4：混合精度演算を安全に導入する

B-3では、forward、backward、optimizerの更新をつなげ、勾配を監視しながらTransformerを学習させる基本を扱いました。学習が正しく進むことを確認できた次に考えたいのが、GPUメモリと計算時間です。モデル・バッチサイズ・系列長を大きくすると、単精度のfloat32だけで計算する構成はすぐに制約へぶつかります。

そこで使われるのが混合精度演算です。計算の一部をfloat16（FP16）またはbfloat16（BF16）で実行し、精度が重要な処理は必要に応じてfloat32で保持します。うまく使えばメモリ使用量や実行時間を改善できますが、「16ビットにすれば常に速く、同じ精度になる」という意味ではありません。GPU、モデル、系列長、演算、実装によって効果も安定性も変わります。

今回は、混合精度を安全に導入するための判断基準と、PyTorchのtorch.autocast・torch.amp.GradScalerを使った最小の訓練ステップを整理します。次回B-5では、型の精度ではなくAttentionのメモリアクセスそのものを最適化するFlashAttentionへ進みます。

1. 混合精度は何を解決するのか？
2. FP16・BF16・FP32はどのように違うのか？
3. Autocastは何を自動化し、何を自動化しないのか？
4. GradScalerはなぜFP16で必要になるのか？
5. PyTorchで安全な訓練ステップを書く
6. FP16とBF16をどのように選ぶのか？
7. 効果はどのように測るのか？
8. よくある失敗と切り分け方
9. まとめ：混合精度を、検証可能な最適化にする
10. 今回のブログの考察：B-5への橋渡し
参考文献
このシリーズの案内

1. 混合精度は何を解決するのか？

GPU上のTransformer学習では、モデル重みだけでなく、勾配、Optimizerの状態、forward時に保存する活性化がメモリを使います。系列長やバッチサイズを増やすと、特に活性化とAttentionまわりの使用量が大きくなります。

混合精度では、すべてを同じ型に固定せず、演算の種類や実行環境に合わせて低精度と高精度を使い分けます。低精度の行列演算をGPUが効率よく処理できる環境では、メモリ使用量と計算時間の両方を改善できる可能性があります。

ただし、効果を事前の割合で約束することはできません。小さなモデルでは型変換やデータ転送のオーバーヘッドが目立つことがあり、対応していないGPUでは期待した高速化を得られない場合もあります。導入の判断は、必ず自分のモデル・バッチ・GPUでの実測に基づけるべきです。

2. FP16・BF16・FP32はどのように違うのか？

浮動小数点数は、表現できる数値の範囲と細かさにトレードオフがあります。FP16とBF16はどちらも16ビットですが、指数部と仮数部の配分が異なります。

型	ビット幅	数値の範囲	相対的な精度	実務上の見方
FP32	32ビット	広い	高い	基準となる単精度。重み・Optimizer状態などで使われることが多い
FP16	16ビット	FP32より狭い	BF16より細かい	対応GPUでは高速な演算を期待できるが、値の大きさによってoverflow・underflowに注意が必要
BF16	16ビット	FP32に近い	FP16より粗い	大きさの範囲を保ちやすく、対応環境ではLLM学習で採用されることが多い

FP16は、非常に小さな勾配が0へ丸め込まれるunderflowを起こすことがあります。そのため、FP16で学習する場合は、通常GradScalerでLossを一時的に拡大してからbackwardを行います。一方、BF16はFP32に近い指数部を持つため、一般にGradScalerを必要としない構成が多いものの、すべてのモデル・GPUで無条件に安定するわけではありません。

型の選択は「BF16が新しいから」ではなく、使用GPUの対応状況、既存モデルの検証結果、メモリ制約で決めます。まずFP32を基準にLossが下がることを確認してから、FP16またはBF16へ切り替える順序が安全です。

3. Autocastは何を自動化し、何を自動化しないのか？

torch.autocastは、指定した範囲の演算について、実行環境と演算の種類に応じた型を選ぶ仕組みです。利用者が線形層や行列積のたびにhalf()やfloat()を呼ぶ必要を減らし、低精度で実行しやすい演算と、より高い精度を保つべき演算を組み合わせます。

ただし、Autocastは学習全体を自動的に安全にする機能ではありません。少なくとも、次の点は利用者が管理します。

Autocastに任せられること	利用者が確認すること
対象範囲内の多くの演算に対する型の選択	GPUがFP16・BF16を適切にサポートしているか
forward内の型変換を必要最小限にすること	Lossが有限値か、勾配にNaN・Infがないか
演算ごとの安全な精度選択を補助すること	FP16時にGradScalerを使うか、勾配クリッピングの順序が正しいか
backwardで使うforward演算の記録を支えること	評価時の精度・生成品質・メモリ・速度をFP32と比較すること

よくある誤解は、モデル全体をmodel.half()へ変換すれば混合精度になるというものです。これはすべてのパラメータと演算をFP16へ寄せる操作であり、Autocastによる選択的な混合精度とは異なります。まずはモデル重みを通常の型で保持し、forwardをAutocastのコンテキストで囲む方法から始めましょう。

4. GradScalerはなぜFP16で必要になるのか？

FP16では表せる数値の範囲が狭いため、非常に小さい勾配が0になり、パラメータ更新が失われることがあります。GradScalerは、backwardの前にLossを拡大し、勾配がFP16で表現できる範囲に入りやすくする仕組みです。

更新の流れは、次の順序で考えると分かりやすくなります。

Autocastの範囲でforwardとLoss計算を行う
ScalerでLossを拡大してbackwardを行う
Optimizerに保持された勾配を元の大きさへ戻す
必要なら勾配クリッピングを行う
Optimizerを更新し、Scalerの係数を更新する

勾配クリッピングを行うなら、必ずunscale_の後に実行します。スケールされたままの勾配ノルムをクリップすると、設定した閾値が本来の意味を持たなくなるためです。B-3で扱った勾配クリッピングと混合精度を組み合わせるとき、ここは特に間違えやすい箇所です。

GradScalerはInfやNaNを検出した場合、その更新をスキップし、スケールを調整します。これは異常を「直す」仕組みではありません。更新が繰り返しスキップされるなら、学習率、データ、損失、モデル実装に問題がないかを調べる必要があります。

5. PyTorchで安全な訓練ステップを書く

以下は、B-1の言語モデルが(batch, sequence, vocab_size)のlogitsを返し、B-2のバッチがinput_idsとlabelsを持つことを想定した最小例です。CUDA環境でのFP16またはBF16を対象とし、FP16の場合だけGradScalerを有効にします。

from dataclasses import dataclass

import torch
import torch.nn.functional as F

@dataclass
class AmpConfig:
    dtype: torch.dtype = torch.bfloat16
    enabled: bool = True
    max_grad_norm: float = 1.0

class MixedPrecisionTrainer:
    def __init__(self, model, optimizer, amp_config: AmpConfig):
        self.model = model
        self.optimizer = optimizer
        self.amp_config = amp_config

        if not torch.cuda.is_available():
            raise RuntimeError("この例はCUDA環境を前提にしています")

        # BF16では通常、GradScalerを使わない。
        use_scaler = (
            amp_config.enabled and amp_config.dtype == torch.float16
        )
        self.scaler = torch.amp.GradScaler("cuda", enabled=use_scaler)

    def train_step(self, batch: dict[str, torch.Tensor]) -> dict[str, float]:
        input_ids = batch["input_ids"].cuda(non_blocking=True)
        labels = batch["labels"].cuda(non_blocking=True)
        self.optimizer.zero_grad(set_to_none=True)

        with torch.autocast(
            device_type="cuda",
            dtype=self.amp_config.dtype,
            enabled=self.amp_config.enabled,
        ):
            logits = self.model(input_ids)
            loss = F.cross_entropy(
                logits[:, :-1, :].reshape(-1, logits.size(-1)),
                labels[:, 1:].reshape(-1),
                ignore_index=-100,
            )

        self.scaler.scale(loss).backward()

        # クリッピング前に、FP16で拡大された勾配を元へ戻す。
        self.scaler.unscale_(self.optimizer)
        grad_norm = torch.nn.utils.clip_grad_norm_(
            self.model.parameters(),
            self.amp_config.max_grad_norm,
        )

        self.scaler.step(self.optimizer)
        self.scaler.update()

        return {
            "loss": loss.detach().float().item(),
            "grad_norm": float(grad_norm),
        }

この例では、入力のlabelsにpadding位置を表す-100が入っている前提です。ignore_index=-100により、その位置はLoss計算から除外されます。paddingの作り方とマスクの扱いはB-2で確認してください。

また、torch.autocastの対象はforwardとLoss計算です。backwardを明示的にAutocastで囲む必要はありません。モデルとOptimizerの初期化、学習率スケジューラ、分散学習との組み合わせは、後続のB-6以降で追加します。

6. FP16とBF16をどのように選ぶのか？

選択に迷ったときは、サポート状況と検証結果から決めます。

条件	最初に検討する選択	確認すること
GPUがBF16に対応し、モデルが安定して学習する	BF16	FP32と比べたLoss、品質、メモリ、スループット
BF16が使えない、または既存環境がFP16中心	FP16 + GradScaler	GradScalerの更新スキップ、勾配ノルム、NaN・Inf
低精度化後に品質や安定性が崩れる	FP32へ戻して原因を切り分ける	データ、学習率、実装、演算ごとの問題
CPUだけで実験する	まずFP32	GPU向けの速度改善を前提にしない

BF16は指数部が広く、FP16よりoverflow・underflowを起こしにくい傾向があります。しかし、対応GPUやカーネル、モデルの設定によって効果は異なります。「BF16ならGradScalerが不要なことが多い」という実務上の指針と、「一切の数値問題が起きない」という保証は分けて考えましょう。

7. 効果はどのように測るのか？

混合精度を導入したら、FP32との比較を同じ条件で行います。モデル、データ、乱数シード、バッチサイズ、系列長、評価データをそろえないと、型の違いによる効果を判断できません。

観点	代表的な記録	判断の例
品質	学習Loss、検証Loss、目的タスクの指標	FP32との差が用途で許容できるか
安定性	NaN・Infの回数、更新スキップ、勾配ノルム	FP16で異常が継続していないか
メモリ	最大GPUメモリ、OOMの有無	同じGPUでバッチや系列長を増やせるか
速度	学習トークン毎秒、1ステップの時間	ウォームアップ後の複数ステップで比較する

GPUメモリは、PyTorchのピーク使用量を記録して比べられます。速度を測るときは、最初の数ステップに含まれるCUDAカーネルの初期化やコンパイルの影響を避けるため、ウォームアップ後に複数回計測します。

torch.cuda.reset_peak_memory_stats()

# ウォームアップ後に、同じ train_step を複数回実行する。
for _ in range(20):
    metrics = trainer.train_step(batch)

peak_memory_gb = torch.cuda.max_memory_allocated() / 1024**3
print(f"loss={metrics['loss']:.4f}")
print(f"peak memory={peak_memory_gb:.2f} GB")

この数値だけで採用を決める必要はありません。メモリが減っても、品質が用途に必要な水準を下回るなら意味がありません。逆に、品質差が小さく、同じGPUで十分大きなバッチや系列長を扱えるなら、混合精度は有効な選択になり得ます。

8. よくある失敗と切り分け方

8.1 FP16へ変えた直後にLossがNaNになる

まずFP32で同じ設定が安定するかを確認します。FP32でも不安定なら、混合精度だけが原因とは言えません。FP16だけで起きる場合は、GradScalerの有効化、学習率、勾配クリッピングの順序、入力や損失の異常値を確認します。

8.2 GradScalerを使いながら、unscale_前に勾配をクリップする

スケールされた勾配でノルムを測ると、クリッピング閾値が実際より厳しく働きます。scale(loss).backward()の後にunscale_(optimizer)を呼び、その後でclip_grad_norm_を実行します。

8.3 精度を変えた後、FP32と違うバッチサイズで比較する

混合精度でメモリが空いたためにバッチサイズを増やすこと自体は有用です。ただし、速度や品質をFP32と比較するときは、同じ条件の結果と、バッチサイズを増やした結果を別の実験として記録します。比較の問いを混ぜないことが重要です。

8.4 model.half()だけで解決しようとする

全てをFP16へ固定すると、数値的に敏感な演算まで低精度で実行される可能性があります。まずはAutocastを使い、必要な範囲で低精度を適用します。特別な事情があって型を固定する場合も、FP32基準と比較して検証します。

9. まとめ：混合精度を、検証可能な最適化にする

今回の要点を整理します。

混合精度は、低精度演算と高精度演算を使い分け、GPUメモリと計算効率の改善を目指す手法である
FP16はGradScalerと組み合わせてunderflowを防ぐことが多く、BF16は広い数値範囲を持つ
Autocastは型選択を補助するが、学習の安定性や評価の責任まで自動化するものではない
GradScaler使用時は、unscale_の後に勾配クリッピングを行う
導入効果は、品質・安定性・メモリ・速度を同じ条件でFP32と比較して判断する
異常が出た場合は、まずFP32基準へ戻し、混合精度以外の原因も含めて切り分ける

混合精度は、モデルを魔法のように高速化する設定ではありません。基準構成を正しく学習させたうえで、同じ品質をより少ない資源で得られるかを確かめるための最適化です。この順序を守ることで、速度を追うあまり学習の正しさを失うことを避けられます。

10. 今回のブログの考察：B-5への橋渡し

B-4では、演算の精度を使い分けて、モデル全体のメモリと計算効率を改善する方法を見ました。しかし、長いコンテキストを扱うTransformerでは、型を低精度にするだけでは解決しない制約があります。Attentionはトークン対の計算とメモリアクセスが大きくなり、系列長が増えるほどボトルネックになりやすいためです。

次回B-5では、Attentionの数学的な結果を変えずに、計算途中のデータをGPUメモリへ何度も書き戻さないFlashAttentionを扱います。混合精度で整えた実行基盤を土台に、Attentionそのものの効率を見直していきましょう。

参考文献

Micikevicius, P., Narang, S., Alben, J., et al. (2018). “Mixed Precision Training.” ICLR 2018. https://arxiv.org/abs/1710.03740
PyTorch. “Automatic Mixed Precision package – torch.amp.” https://pytorch.org/docs/stable/amp.html
PyTorch. “torch.autocast.” https://pytorch.org/docs/stable/amp.html#torch.autocast
PyTorch. “torch.amp.GradScaler.” https://pytorch.org/docs/stable/amp.html#gradient-scaling

このシリーズの案内

ブログB「Transformer実装」では、Aシリーズで整理した設計を、PyTorchの実装・最適化・運用上の工夫へ段階的に落とし込みます。

#	テーマ	内容
B-1	Transformerブロック実装	Causal Attention、FFN、Pre-Norm、次トークン予測
B-2	バッチ処理と位置エンコーディング	padding、DataLoader、RoPE
B-3	自動微分と最適化	backward、AdamW、勾配クリッピング
B-4	混合精度演算	FP16・BF16、Autocast、GradScaler（今回）
B-5	FlashAttention	Attentionの高速化とメモリ効率
B-6	分散学習とDDP	複数GPUでの学習
B-7	メモリ最適化	活性化・KVキャッシュ・チェックポイント

B-3：自動微分とバックプロパゲーション最適化

次に読む

B-5：FlashAttentionと高効率Attention実装

実装を深掘りしたい方へ

B-2：バッチ処理と位置エンコーディングの実装
B-6：分散学習とDDP

LLM【事前学習Ⅱ：基礎理論 B-4】