Transformersコミュニティ、Mamba2の正規化におけるグループ分割の欠落を報告　非融合パスが学習時の中核計算と異なる可能性

Transformersコミュニティは10月12日、Mamba2の計算整合性に関する問題を報告しました。n_groups > 1の場合、非融合のゲート付きRMSNormは中間次元全体で平均二乗値を計算しますが、上流の融合学習パスではSSMグループごとに正規化しています。報告者は指定された開発コミットを調査し、Transformers 5.5.0で最小再現を実行しました。現時点でIssueは未解決で、関連するPRはありません。公開された報告

違いは正規化の範囲にあります。現行のMambaRMSNormGatedは最終次元に沿って二乗平均を計算し、インスタンスの構築時にグループ情報を受け取りません。一方、融合パスにはngroupsが渡されます。グループごとに値のスケールが異なる場合、全体平均を使うと、一方のグループの入力がもう一方の出力に影響します。報告ではFP32の小型テンソルを使って公式の参照実装と比較し、出力の最大差は約0.802でした。また、グループごとの参照実装ではゼロになるはずのグループ間導関数が、現行実装ではゼロにならないことも測定されました。これらは報告者による測定であり、独立した再現はまだ行われていません。再現コードと結果

この検証は実運用にも関わります。公式ドキュメントではn_groupsのデフォルト値を8とし、Codestral Mambaが8グループを使うこと、またPyTorchとCUDAカーネルの2種類のフォワードパスがあることを説明しています。ドキュメントは両者にわずかな数値差が生じる可能性を示していますが、今回報告された計算上の分析によれば、グループ範囲の違いは関数の定義に関わります。エンジニアはこれを演算順序や低精度に起因する誤差と分けて検証する必要があります。Mamba2公式ドキュメント

報告者は、テンソルをグループ単位に再形成して各グループでRMSNormを計算する修正案を提示しています。この方法では既存の重みの形状を維持でき、報告者による局所的な出力と勾配の比較では問題が解消されました。ただし、融合SSM操作、事前学習済みモデル、生成品質を含む検証はまだ行われていません。また、Codestralのゲート処理順序は今回の再現範囲外のため、品質への影響が確認されたとは言えません。今後はメンテナーによるレビューに加え、融合学習、通常のフォワード、キャッシュを使ったデコードの間で、モデル全体の整合性が検証されるかが焦点です。検証範囲と修正案