エラー解決

【PyTorch】Mixed Precision Training (AMP) のエラー解決とベストプラクティス：メモリ削減と高速化を実現

· 15分で読めます

問題の概要：Mixed Precision Training (AMP) の導入で発生する典型的なエラー

PyTorchのAutomatic Mixed Precision (AMP) は、FP16（半精度浮動小数点）とFP32（単精度）を自動的に使い分けることで、GPUメモリ使用量を削減し、訓練速度を向上させる強力な機能です。しかし、導入時には以下のようなエラーや予期せぬ動作に遭遇することが多く、特に初心者から中級者の開発者を悩ませます。

  • 「NaN」や「Inf」の損失値: 訓練中に損失値が突然「nan」や「inf」になり、学習が破綻する。
  • 精度の大幅な低下: FP32での訓練と比較して、モデルの最終精度（Accuracyなど）が大きく低下する。
  • エラーメッセージ: UserWarning: スケーラーがスキップされました。損失スケールが0です。 や RuntimeError: value cannot be converted to type half without overflow: などの警告・エラー。
  • メモリ削減効果が小さい: 理論通りにメモリ使用量が減らず、バッチサイズを大きくできない。

これらの問題は、AMPの仕組みを理解せずに単純にコードに追加した場合に頻発します。

原因の解説：なぜエラーが起こるのか？

AMPの核心は、FP16の「表現範囲の狭さ」にあります。FP32に比べ、FP16は表現できる数値の範囲（ダイナミックレンジ）が極端に狭いです。この特性が以下の問題を引き起こします。

1. アンダーフローとオーバーフロー

訓練中の勾配は、非常に小さな値（1e-8など）になることがあります。FP16ではこのような小さな値が「0」（アンダーフロー）として表現され、勾配情報が消失します。逆に、大きな値は「無限大（inf）」（オーバーフロー）に丸められ、nan損失の原因となります。

2. 不適切なスケーリング

AMPは「勾配スケーラー（GradScaler）」を使用してこの問題を緩和します。スケーラーは損失を適切に拡大（スケールアップ）してFP16での計算を行い、その後勾配を縮小（スケールダウン）してFP32のパラメータを更新します。このスケーリングの設定（成長係数やバックオフ期間）が不適切だと、効果が得られなかったり、逆に不安定になります。

3. FP16非対応の演算

一部の演算（ソフトマックス、レイヤーノーマライゼーションの分散計算など）は、FP16では数値的不安定性を引き起こす可能性があります。AMPはオペレータのホワイトリスト/ブラックリストに基づき、これらの演算を自動的にFP32で実行しますが、カスタムレイヤーや複雑な操作ではこの自動判定が機能しない場合があります。

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

以下に、安定したAMP訓練を実現するための具体的な手順を示します。

ステップ1: 基本的なAMPコードの実装

まずは、公式推奨の基本的な実装パターンを確実に導入します。

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

# 訓練ループの開始前にスケーラーを初期化
scaler = GradScaler()

for epoch in range(num_epochs):
    for data, target in train_loader:
        data, target = data.cuda(), target.cuda()
        optimizer.zero_grad()

        # 順伝搬をautocastコンテキスト内で実行
        with autocast():
            output = model(data)
            loss = criterion(output, target)

        # スケーラーを使って逆伝搬と最適化を実行
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

ステップ2: 勾配スケーラーの詳細設定（最重要）

デフォルト設定では不安定な場合、スケーラーのパラメータを調整します。特にgrowth_intervalを調整することが有効です。

# スケーラーの詳細設定例
scaler = GradScaler(
    init_scale=65536.0,     # 初期スケール値 (2^16)
    growth_factor=2.0,      # オーバーフローがなければスケールを2倍にする
    backoff_factor=0.5,     # オーバーフローが発生したらスケールを0.5倍にする
    growth_interval=2000,   # growth_intervalステップ連続でオーバーフローがなければスケールを増加
    enabled=True
)

growth_intervalを大きくする（デフォルトは2000）と、スケールの増加が保守的になり、安定性が向上します。逆に不安定性が続く場合は、init_scaleを小さく（例: 1024.0）することも検討します。

ステップ3: FP32で実行すべき演算の強制指定

数値的に敏感なカスタム関数や、損失関数の一部には明示的にFP32を使用させます。

# 例1: 損失関数内の特定計算をFP32で
def custom_loss(output, target):
    with autocast(enabled=False): # このブロック内はFP32強制
        sensitive_part = some_sensitive_operation(output)
    # その他の計算はAMPに任せる
    return criterion(sensitive_part, target)

# 例2: モデルの特定モジュールをFP32のままにする（初期化時）
class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(...) # AMP対象
        # 最終層はFP32で固定（数値的に敏感な場合）
        self.final_layer = nn.Linear(256, 10).float()

ステップ4: 勾配クリッピングの併用

AMP使用時は、スケーラーを考慮した勾配クリッピングが必須です。従来のtorch.nn.utils.clip_grad_norm_は使えません。

scaler.scale(loss).backward()
# スケーラーを考慮した勾配クリッピング
scaler.unscale_(optimizer) # まず勾配を「アンスケール」
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) # クリップ
scaler.step(optimizer)
scaler.update()

警告: scaler.unscale_は、scaler.stepの前に1回だけ呼び出す必要があります。複数回呼び出すとエラーになります。

ステップ5: ナン・無限大の監視とデバッグ

訓練中に問題が発生した場合、以下のコードでデバッグします。

# 損失がnan/infになっていないかチェック
if not torch.isfinite(loss):
    print("非有限の損失を検出しました。スケーラー状態をリセットします。")
    scaler.update(new_scale=scaler.get_scale() * backoff_factor) # 手動でスケールダウン

# 勾配にnan/infがないかチェック（デバッグ用）
for name, param in model.named_parameters():
    if param.grad is not None:
        if torch.isnan(param.grad).any() or torch.isinf(param.grad).any():
            print(f"勾配に問題あり: {name}")

コード例・コマンド例：完全な訓練ループ

上記のベストプラクティスを全て組み込んだ、実用的な訓練ループの例です。

def train_one_epoch(model, train_loader, criterion, optimizer, scheduler, scaler, epoch):
    model.train()
    total_loss = 0

    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.cuda(non_blocking=True), target.cuda(non_blocking=True)

        optimizer.zero_grad(set_to_none=True) # メモリ効率向上

        # 順伝搬 (AMP有効)
        with autocast():
            output = model(data)
            loss = criterion(output, target)

        # 逆伝搬と最適化
        scaler.scale(loss).backward()

        # 勾配クリッピング
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

        scaler.step(optimizer)
        scaler.update()
        scheduler.step()

        # 損失の記録（スケール前の損失を使用）
        total_loss += loss.item()

        # 200ステップごとにスケーラー状態をログ出力
        if batch_idx % 200 == 0:
            print(f'Epoch: {epoch} | Step: {batch_idx} | Loss: {loss.item():.4f} | Scale: {scaler.get_scale():.4f}')

    return total_loss / len(train_loader)

# 初期化
scaler = GradScaler(init_scale=65536.0, growth_interval=2000)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs)

# 訓練実行
for epoch in range(num_epochs):
    avg_loss = train_one_epoch(model, train_loader, criterion, optimizer, scheduler, scaler, epoch)
    print(f'[Epoch {epoch}] Average Loss: {avg_loss:.4f}')

まとめ・補足情報

PyTorch AMPは「魔法の機能」ではなく、その仕組みを理解した上で適切に設定する必要があります。安定した訓練のためには、勾配スケーラーの詳細設定と勾配クリッピングの適切な実装が最も重要です。

追加のベストプラクティス:

  • バッチサイズ: AMP導入でメモリが節約できるので、バッチサイズを2倍にしてみましょう。これがAMPの最大のメリットです。
  • パフォーマンス比較: 必ずFP32でのベースラインと、AMP使用時の最終精度・訓練時間・最大バッチサイズを比較検証してください。
  • PyTorchバージョン: 常に最新の安定版PyTorchを使用してください。AMP関連の改善は活発に行われています。
  • NVIDIA Tensor Cores: Voltaアーキテクチャ以降のGPU（V100, A100, RTXシリーズ等）で最大の効果を発揮します。古いGPUでもメモリ削減効果は得られます。

最初はエラーに遭遇するかもしれませんが、本記事で紹介したステップに従って設定を見直せば、メモリ効率と訓練速度を大幅に向上させる「Mixed Precision Training」を確実に活用できるようになるでしょう。

この記事は役に立ちましたか？