PyTorchコミュニティがBF16 Dropoutの乱数列の重複を報告、連続呼び出しでマスクが相関する可能性

PyTorchは10月28日、学習時のランダム性に関する不具合報告を受け付けた。NVIDIA L40上でeager modeを使い、aten::native_dropoutを連続して呼び出すと、特定のBF16テンソルから生成される2つのマスクに規則的な重複が生じるという。報告者は2.7.0+cu126と2.14.1+cu126の両方で同じ現象を確認した。これは2つのバージョンでの実測結果であり、その間のすべてのバージョンに影響すると断定する根拠にはならない。不具合報告

再現条件は、形状が(4096, 128)で、連続配置され、メモリアドレスが16バイト境界にアラインされたBF16入力と、ドロップアウト確率0.5。マスクを8要素ずつに分けると、1回目の呼び出しの後半4要素が、2回目の呼び出しの先頭4要素と完全に一致する。独立したベルヌーイ分布のマスクで期待される一致率は約半分だ。テストでは、2回の呼び出しの前に一度だけシードを設定した。FP32に変更するか、BF16の開始アドレスを8バイトずらすと、この特定の重複は発生しなくなる。再現結果と比較

報告者は、ベクトル化カーネルが消費する乱数の量と、Philoxカウンターで予約される量が一致していない可能性を指摘している。公開されている2.14.1のソースコードでは、乱数の生成回数はベクトル幅に応じて調整される一方、カウンターのオフセットはUNROLLを使って計算される。1回の処理で8要素を扱う際に、乱数の位置を4つ分しか予約しない場合、次の呼び出しで前回の乱数列が再利用される可能性がある。この説明は報告された結果と整合するが、根本原因の分析はまだ確認されていない。CUDAカーネルのソースコード

技術的に重要なのは、マスクの独立性だ。公式ドキュメントによると、Dropoutは学習時にベルヌーイ分布に従って要素をゼロにし、各forward呼び出しで独立にサンプリングすることで、ニューロンの共適応を抑制する。そのため、呼び出しをまたいだ固定的な相関は正則化の挙動を変える可能性がある。ただし、完全なモデルでの収束や精度に関するデータはなく、学習への影響を定量化することはできない。Dropoutのドキュメント

確認時点で、このissueはオープンのままで、自動分類ラベルのみが付いており、メンテナーによる確認や関連する修正は見当たらない。エンジニアリングチームは最小再現コードを使い、自社環境のハードウェア、テンソル形状、実際の実行経路を確認できる。今後はGPUをまたいだ検証、乱数オフセットの修正、安定版ブランチへのバックポートを追跡する必要がある。また、結果をコンパイルモードや他の融合カーネルにそのまま当てはめることはできない。対応状況