PyTorchコミュニティ、大規模 bmm の範囲外書き込みを報告――GB200で再現、安定版への修正適用に差

PyTorchコミュニティは10月27日、最小再現コードを添えたバグ報告を提出しました。NVIDIA GB200、PyTorch 2.13.0+cu130で、縮約次元K=1のtorch.bmmが範囲外書き込みを引き起こす可能性があるというものです。確認時点でissueはオープンのままで、メンテナーによるバックポート予定の確認はありません。報告者によると、記事の整理にはAIを利用しましたが、実験とブランチの確認は本人が行いました。バグ報告

bmmは2つの3次元テンソルをバッチごとに乗算します。入力形状が[b,n,m]と[b,m,p]の場合、出力形状は[b,n,p]です。中間次元が1なら、各バッチの演算は外積に相当するため、入力が小さくても出力が急速に大きくなることがあります。大規模テンソルのインデックスを調べる際、入力サイズだけを見てはいけない理由の一つです。PyTorch APIドキュメント

報告されたテストでは、入力を[512,B,1]と[512,1,512]に設定しています。Bを8208から8209に増やすと、最後のバッチの開始要素のオフセットがINT32_MAXを超え、ネイティブカーネルでエラーが発生しました。報告者は、オフセットの乗算にint32が使われており、オーバーフローによって誤った位置に書き込まれることを根本原因として挙げています。比較テストでは、TORCH_DISABLE_NATIVE_JIT=1を設定してcuBLASに切り替える方法と、関連するインデックスをint64に変換する方法のどちらも、提示されたケースを通過しました。再現コードと境界テスト

この種のエラーは、その後の診断を妨げることがあります。NVIDIAのドキュメントでは、cudaErrorIllegalAddressはカーネルが無効なアドレスにアクセスしたことを示し、プロセスを不整合な状態にするため、その後のCUDA処理でもエラーが続き、プロセスを終了して再起動する必要があると説明されています。したがって、エンジニアリングチームは境界ケースを別々に実行し、連鎖的な失敗を複数の独立した問題と誤認しないようにする必要があります。CUDAエラー定義

影響するバージョン範囲は、引き続き明確に区別する必要があります。実機での再現は2.13.0のGB200環境で行われ、2.14のリスクは報告者によるブランチコードの比較に基づいています。mainブランチと2.15ブランチで修正済みという説明も、同報告によるものです。今後は、メンテナーがstable版への影響を確認するか、バックポートを予定するか、また代替実行経路が実際のワークロードに与える性能影響を追跡する必要があります。この報告から、すべてのGPUやすべての行列乗算が影響を受けると推定することはできません。バージョン確認とバックポート要請