PyTorch でモデル学習中に発生する NaN 値は、エラーを発生させずに静かにモデルを破損させるため、従来からのデバッグ手法である torch.autograd.set_detect_anomaly には致命的な欠点があります。この機能は計算グラフ全体を保持して逆伝播を検査するため、GPU では実行速度が 50 倍から 100 倍、CPU でも 10 倍以上遅くなり、さらに NaN が発生した本当の層ではなく、勾配計算が破綻した後の層を報告する傾向があります。これを解決するために、著者は順伝播フックを用いた軽量な検出ツール「PyTorch NaN Detector」を開発しました。このツールは各層の出力に直接チェックフックを装着し、NaN または無限大を即座に検出する仕組みです。基準となるメソッドと比較して、このアプローチは 1 回の実行あたり約 3 秒のオーバーヘッドのみで済み、大規模モデルにおける学習時間の増加を劇的に抑えています。設計の鍵となるのは、Tensor のチェック自体はロックなしで実行し、共有状態の変更のみをスレッドセーフなロックで保護する点です。これにより、マルチワーカー環境でも競合状態を防ぎながら高速な監視を実現しています。また、学習率が高すぎて勾配ノルムが爆発するといった根本原因の兆候を、順伝播のチェックよりも前の段階で捉えるため、勾配ノルムガード機能を備えています。これにより、NaN が活性化値に伝播する前にバッチ 1 ステップ目で警告を発することが可能です。ユーザーは、特定のパラメータや Dropout など、正常に非有限値を出力する可能性がある層を除外設定したり、バッチインデックスや層名、統計値を含む構造化されたイベントログを取得したりできます。このツールは単に NaN の発生場所を特定するだけでなく、なぜ発生したのかを推測するための重要なデータを提供し、ミックス精度学習におけるスケールオーバーフローなどの問題調査を支援します。ただし、これは学習時の衛生状態を保つための代替手段ではなく、勾配クリッピングや適切な学習率スケジューリングといった基本的なベストプラクティスを補完するデバッグツールとして位置づけられています。ソースコードとベンチマークデータは GitHub 上で公開されており、MIT ライセンスで利用可能です。