Transformers の報告：Gemma 2 の注意力ソフトキャッピングが適用されず、バックエンドの切り替えでモデル計算が変わる可能性

Hugging Face Transformersで10月21日、推論の再現性に関する報告があった。Gemma 2の設定にある注意力ソフトキャッピングが、注意力バックエンドを切り替えると適用されない可能性があり、実行時に警告も表示されないという。報告された環境はTransformers 5.20.2、PyTorch 2.13.0、H100 1基で、再現コードも提供されている。問題報告

Gemma 2のこの処理では、注意力スコアが後続の計算に進む前に、50 × tanh(x / 50)によって大きな値が圧縮される。公式設定ドキュメントでは、attn_logit_softcappingのデフォルト値は50とされている。また、最終出力logitsにも独立したソフトキャッピングがあるため、両者は分けて確認する必要がある。今回の報告が対象としているのは前者だ。Gemma 2 ドキュメント

報告者の分析によると、eagerとflex_attentionではこの変換が適用される一方、Gemma 2がデフォルトで使用するsdpaは渡されたパラメーターを無視する。Flash Attentionの経路では、カーネルがsoftcapを受け付けるかどうかによって挙動が異なる。報告では、ROCmのkernels-community/aiter-flash-attnはこのパラメーターを受け付けないと指摘されている。そのため、すべてのFlash Attentionカーネルが影響を受けると一般化することはできない。バックエンド分析

報告ではGemma 2 2B、WikiTextの検証データ、FP32を使って比較している。2,048 tokenの単一セグメントでは損失値の差は小さかったが、eagerとSDPAの出力logitsの最大絶対差は約2.443だった。ソフトキャッピングを無効にすると、表示精度ではeagerとSDPAの差はゼロになった。この結果は、バックエンドで処理が省略されているとの判断を裏付けるものだが、下流タスク全般の性能低下を示すには不十分だ。再現結果

エンジニアリングチームにとって、これは高速化バックエンドを切り替える際、スループットやメモリだけでなく、モデルの計算が完全に維持されているかも確認すべきことを意味する。評価記録にはバックエンド、カーネル、バージョンを残し、固定入力で数値を比較するのが望ましい。今後は、上流で能力チェック、適切な警告、回帰テストが追加されるかを追う必要がある。現時点で問題は未解決のままで、正式な修正バージョンはまだ確認されていない。