Transformersコミュニティ、全層がスライディングウィンドウのモデル生成時にクラッシュする問題へのキャッシュオフロード修正を提案

Transformersコミュニティは10月24日、特定のモデル設定でoffloaded_staticを使うと生成が直ちに中断する問題に対処するPR #49955を提出した。提出者によると、すべての注意層がスライディングウィンドウ層の場合、初回のフォワードパスでValueErrorが発生する可能性がある。影響を受ける設定にはPhi-3-miniとMistral-7B-v0.1が含まれる。これは新たに提案された修正であり、正式リリース済みの修正ではない。修正提案

キャッシュオフロードは、KV状態の一部をCPUに移してGPUメモリ使用量を抑え、計算前にプリフェッチする仕組みだ。公式ドキュメントによると、静的キャッシュは既定でスライディングウィンドウ層以外のみをオフロードする。スライディングウィンドウキャッシュのサイズには上限があり、転送コストによってメモリ節約の利点が相殺される可能性があるためだ。この既定値は動的キャッシュとは異なるため、「オフロードが有効」という表示だけを見て、すべての層がCPUに移されると判断することはできない。キャッシュ戦略のドキュメント

提案の分析によると、全層がスライディングウィンドウの構成では、オフロード対象層のリストが空になる一方、プリフェッチ関数は次の対象を探そうとして例外が発生する。修正では、オフロードが必要な層がない場合にそのまま戻るようにする。つまり修正後も、この構成のスライディングウィンドウキャッシュはアクセラレーター上に残る。「クラッシュしなくなる」ことを、メモリ節約効果が追加で得られるという意味に解釈してはならない。原因と変更の説明

提出者は、縮小したPhi-3設定と小型のMistralモデルをRTX 2060 SUPERで実行し、修正後の出力が通常の静的キャッシュと一致することを示した。これは特定の経路が正しく動作することを裏付けるが、完全なモデルのデプロイ時の性能を示すには不十分だ。VRAMの少ないGPUを使うチームは、今後のメンテナーによるレビューとリリースへの取り込みを注視し、実際のモデルで生成出力、ピーク時のVRAM使用量、スループットを比較するとよい。現時点では、再現証拠を伴うコミュニティの修正提案として扱うべきだ。再現例とテスト結果