TransformersコミュニティがMiniMaxの精度修正を提案、bf16の再帰的減衰で長系列デコードがずれる可能性

Transformersコミュニティは10月4日、MiniMax Lightning Attentionの低精度演算により、デコード長が伸びるにつれて数値誤差が生じる可能性があると報告しました。事例では開発ブランチ、PyTorch 2.11、CPUを使用しており、問題はトークンごとに更新される減衰係数と再帰状態にあるとされています。修正案は提出済みですが、上流でのレビュー待ちです。問題報告

この演算は、長いコンテキストを扱うモデルにとって特に重要です。公式ドキュメントによると、MiniMax-Text-01はLightning Attention、Softmax Attention、MoEを組み合わせており、Lightning Attentionを7層重ねるごとにSoftmax Attentionを1層配置しています。再帰状態の減衰は古い情報をどの程度保持するかを制御するため、デプロイ時には重みの精度だけでなく、状態更新の計算精度も確認する必要があります。アーキテクチャのドキュメント

報告によると、bf16は1に近い値の精度が不足しており、0.998以上の減衰係数が1に丸められて、該当する注意ヘッドで減衰が止まる可能性があります。報告者は公開設定を使った推定で、Lightning Attentionの4,480個のヘッドのうち388個がこの条件に当てはまるとしています。ただし、これは係数の分析であり、実タスクでの誤り率や品質低下を直接示すものではありません。再現結果と制約

修正案では、減衰に関わる演算とキャッシュ状態をfp32に引き上げ、出力射影の前に元の精度へ戻します。提案者は小型モデルでの再現実験を用い、3,000トークン時のfp32に対する隠れ状態の誤差が約7.39%から0.39%に低下したと報告しました。この結果は修正方針を支持しますが、MiniMaxの完全なモデルを使った長文評価ではありません。修正提案

エンジニアリング上のコストとして、線形注意の状態をfp32にすると、その部分のキャッシュ容量はbf16の場合の約2倍になります。モデル全体のメモリ使用量が2倍になるわけではありません。今後、デプロイチームは正式リリースへの影響範囲、実際の重みを使ったGPUでの検証、修正後のスループットとメモリコストを確認する必要があります。また、短い出力のテストでは累積誤差が見えにくいため、受け入れ評価ではプロンプトのプリフィルとトークン単位のデコードを分けて検証すべきです。報告者はMiniMaxの完全なチェックポイントではまだ実行していないと明記しており、現時点では再現コードを伴う数値問題の報告として捉える必要があります。検証範囲