【PyTorch】Mixed Precision Training (AMP) のエラー解決と実践的ベストプラクティス

問題の概要：AMP使用時の典型的なエラーと課題

PyTorchのAutomatic Mixed Precision (AMP)は、メモリ使用量を削減し、訓練速度を向上させる強力なツールです。しかし、特に初心者・中級者が導入する際には、いくつかの典型的なエラーに遭遇します。主な問題として、「RuntimeError: value cannot be converted to type at::Half without overflow」や「Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!」といったエラーが発生し、訓練プロセスが突然停止することがあります。また、数値的不安定性（NaN損失の発生）や、思ったほど速度が向上しないといったパフォーマンス上の課題もよく報告されます。

原因の解説：なぜこれらのエラーが発生するのか？

AMPのエラーは、主に以下の3つの根本的原因に起因しています。

1. データ型の不適合とオーバーフロー

AMPは、計算グラフの一部を16ビット浮動小数点数（float16/bfloat16）に自動的にキャストします。しかし、非常に大きい値や小さい値（例えば、大きな損失値、softmax前の大きなロジット）は、float16の限られたダイナミックレンジ（約 5.96e-08 ~ 65504）では表現できず、無限大（inf）やゼロ（0）にアンダーフロー/オーバーフローしてしまいます。これが「overflow」エラーの直接の原因です。

2. デバイス不一致エラー

AMPコンテキスト（autocast）内で、CUDA（GPU）テンソルとCPUテンソルが混在する操作を行うと、デバイス不一致エラーが発生します。autocastは主にCUDA演算を最適化するため、意図せずCPUにデータが残っていると問題が起こります。

3. 勾配スケーリングの不適切な管理

GradScalerは、float16でアンダーフローする可能性のある小さな勾配をスケールアップし、オプティマイザのステップ後にスケールダウンします。このスケーリングのタイミング（scaler.step(optimizer)とscaler.update()）を誤ると、勾配が正しく更新されなかったり、スケーリングが累積されてオーバーフローを引き起こしたりします。

解決方法：ステップバイステップのベストプラクティス

ステップ1：安全なAMP訓練環境の構築

まず、基本的で堅牢なAMP訓練ループを構築します。これにより、多くの初歩的なエラーを防げます。

import torch
from torch.cuda.amp import autocast, GradScaler

# 1. デバイス設定（明示的にCUDAを使用）
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

# 2. GradScalerのインスタンス化（推奨設定から始める）
scaler = GradScaler(enabled=True) # enabledフラグで簡単にON/OFF可能

# 3. オプティマイザ（AMPと互換性のあるものを選択）
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

for epoch in range(num_epochs):
    for data, target in dataloader:
        # 4. データを確実に正しいデバイスへ
        data, target = data.to(device), target.to(device)

        optimizer.zero_grad()

        # 5. autocastコンテキスト内でフォワードパスを実行
        with autocast(enabled=True, dtype=torch.float16): # または torch.bfloat16
            output = model(data)
            loss = loss_fn(output, target)

        # 6. スケーラーを使って逆伝播と最適化
        scaler.scale(loss).backward() # loss.backward() ではない！
        scaler.step(optimizer)
        scaler.update() # スケール係数を更新

        # 7. （オプション）スケーラーの状態を定期的に確認
        # print(scaler.get_scale()) # スケール係数を確認

ステップ2：数値的不安定性（NaN/Inf）への対処

損失がNaNになる場合は、以下の対策を講じます。

scaler = GradScaler(init_scale=65536.0, # デフォルトは65536.0。小さくすると安定化するが、アンダーフローリスク増
                    growth_factor=2.0,
                    backoff_factor=0.5,
                    growth_interval=2000)

# 訓練ループ内でNaNを監視
with autocast(dtype=torch.float16):
    output = model(data)
    loss = loss_fn(output, target)

# スケーラーのステップ前に、勾配の状態をチェック（デバッグ用）
# scaler.unscale_(optimizer) # 通常はscale(loss).backward()内で行われる
# torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) # 勾配クリッピングを追加

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

# スケーラーの更新結果を確認。NaNが検出されるとスケールが減少する。
if scaler.get_scale() < 1.0:
    print(f"警告: スケーラーがNaNを検出し、スケール係数を{scaler.get_scale()}に減らしました。")

代替案： float16ではなく、ダイナミックレンジが広いbfloat16を使用する（対応するGPUが必要）。
with autocast(enabled=True, dtype=torch.bfloat16):

ステップ3：モデル固有のAMP無効化（安全策）

特定のレイヤーや関数がfloat16で不安定な場合、その部分だけをfloat32で計算させます。

# 方法A: モデルの特定モジュールをFP32にキャスト
model.embedding_layer = model.embedding_layer.float()
model.some_critical_layer = model.some_critical_layer.float()

# 方法B: autocastコンテキスト内で、特定の関数のみFP32で実行（推奨）
with autocast(dtype=torch.float16):
    # ほとんどの計算はFP16
    x = model.conv_layers(data)
    # この関数だけはFP32で強制実行
    with torch.cuda.amp.autocast(enabled=False):
        x = model.unstable_operation(x.float()) # .float()で明示的にFP32化
    output = model.rest_of_layers(x)

ステップ4：パフォーマンス最大化のための微調整

安定性が確保されたら、速度を最大化します。

# 1. カスタムオプティマイザや勾兵蓄積との併用
accumulation_steps = 4
scaler = GradScaler()

for i, (data, target) in enumerate(dataloader):
    data, target = data.to(device), target.to(device)
    with autocast(dtype=torch.float16):
        output = model(data)
        loss = loss_fn(output, target) / accumulation_steps # 損失を正規化

    scaler.scale(loss).backward()

    if (i + 1) % accumulation_steps == 0:
        # 勾配蓄積ステップごとにオプティマイザを更新
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()

# 2. データローダーのピンメモリと非ブロッキング転送を有効化
dataloader = DataLoader(dataset, batch_size=32, pin_memory=True, num_workers=4)
# 転送時
data = data.to(device, non_blocking=True)

コード例・コマンド例：よくあるエラーとその修正

エラー例1: デバイス不一致

# エラーが発生するコード
with autocast():
    output = model(data.cuda()) # dataはCUDA
    loss = loss_fn(output, target) # targetはCPUに残っている！

# 修正コード
data, target = data.to(device), target.to(device) # 両方を同じデバイスに明示的に移動
with autocast():
    output = model(data)
    loss = loss_fn(output, target)

エラー例2: スケーラーの誤った使用

# 誤ったコード（scaler.update()の位置）
scaler.scale(loss).backward()
scaler.step(optimizer)
optimizer.zero_grad() # ここでゼロ化してしまうと...
scaler.update() # ...updateが呼ばれる前に勾配が消える（問題ないが順序が非標準）

# 推奨される順序（PyTorch公式）
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad() # update()の後にゼロ化

まとめ・補足情報

PyTorch AMPを効果的かつ安全に使用するためのベストプラクティスは、「段階的な導入」と「予防的対策」に集約されます。まずは基本的な訓練ループでAMPを有効化し、安定性を確認してください。数値的不安定性が発生した場合は、GradScalerのパラメータ調整、bfloat16への切り替え、不安定な演算のFP32強制といった対策を講じます。パフォーマンスについては、データパイプラインの最適化（pin_memory, non_blocking）や勾配蓄積との組み合わせでさらに向上が期待できます。

最終チェックリスト：

1. すべてのテンソルが訓練前に正しいデバイス（GPU）に移動しているか？
2. autocastコンテキスト内でフォワードパスと損失計算を行っているか？
3. loss.backward()ではなくscaler.scale(loss).backward()を呼んでいるか？
4. scaler.step(optimizer)とscaler.update()の呼び出し順序は正しいか？
5. NaNが頻発する場合、モデルの初期化や損失関数に問題はないか？（AMP以外の原因も調査）

AMPは適切に使用すれば、リソース制約のある環境や大規模モデル開発において、開発効率を大幅に高めることができる必須の技術です。本記事のプラクティスを参考に、安全かつ高速な訓練を実現してください。