Transformers、FlashAttentionの修正を提案　Ministral3とMistral4のパック済み系列の分離を復元

Transformersコミュニティは11月6日、PR #50450を提出し、FlashAttention使用時にMinistral3とMistral4がパディングなしのパック済み系列を処理する際の引数受け渡しを修正することを提案した。作者の説明によると、両モデルのアテンションモジュールはposition_idsを受け取るものの、それをアテンション関数に渡していなかった。そのためバックエンドが系列の境界を認識できず、文書間で境界をまたいだアテンションが発生する可能性がある。確認時点で、この提案はオープンのままだった。修正提案

この問題は、ファインチューニングでよく使われる効率化手法に関係する。通常のバッチ処理では、短いサンプルを最も長いサンプルと同じ長さまでパディングする。パディングなしの処理では、サンプルを連結して無駄なトークン計算を減らす。Hugging Faceの技術記事によると、FlashAttention 2のこの経路では、サンプル間の境界を保つために位置情報を使い、可変長アテンションで各系列を処理する。境界情報が届かなければ、連結されたデータが相互にアテンション可能な内容として扱われるおそれがある。公式技術解説

作者によると、モデルがクエリのスケーリング用に位置引数を明示的に受け取った後、それをattention_interfaceに渡していなかったことが原因だ。提案ではこの受け渡しを追加し、アテンション関数が実際に引数を受け取ることを確認するテストも加えている。また作者は、既存のパディングなしテストでは、以前の変更後にこの問題を検出できなくなっていたと指摘した。公開ページでは、CIの実行に対するメンテナーの承認待ちと表示されている。したがって、正式なパッケージで修正済みとは判断できない。問題の特定と検証状況

エンジニアリングチームが注意すべき点は、学習データにおけるサンプル間の独立性だ。上記の仕組みから推測すると、本来隔離されるべきサンプル同士が互いの文脈を参照した場合、損失値や収束結果が想定から外れる可能性がある。現時点で提案には学習品質の包括的な比較はなく、影響を受けるバージョンも特定されていない。TRLのドキュメントでは、パッキングとパディングなしの処理は別々のメモリ最適化手法として説明されている。実際の発生条件は、データコレーターとアテンション設定を照合する必要がある。TRLドキュメント

該当モデルを使ってファインチューニングを行うチームは、サンプルを個別に実行した場合と連結して実行した場合の出力を優先的に比較し、先行する文書が後続の文書に影響していないか確認するとよい。今後はメンテナーのレビュー、実際のFlashAttention回帰テストの結果、修正が含まれるバージョンを確認する必要がある。現時点の証拠では、既存の学習タスクへの影響の程度を定量化できない。