SparseBalance: 動的スパース注意によるロードバランスされたロングコンテキスト学習

要旨: 疎（スパース）アテンションは長文コンテキストLLMの学習における計算上のボトルネックを緩和しますが、その分散学習プロセスは、 extit{1)} シーケンス長と extit{2)} 疎性（スパース）感度の両面において極端な異質性を示します。その結果、深刻な負荷分散（バランス）問題と、最適でないモデル精度につながります。既存のアルゴリズムや学習フレームワークは通常、単一の課題に焦点を当てており、これら2つの問題を体系的に同時最適化できていません。そこで本研究では、スパース性とシーケンスの異質性を活用し、モデル精度とシステム効率を共同で最適化する、新しいアルゴリズム—システム共同設計フレームワークである SparseBalance を提案します。まず、ワークロードに応じた動的疎性チューニングを提案し、双方向の疎性調整を用いてストラグラーを解消し、内在するバブルを「無料」で精度に活用します。次に、動的疎性チューニングを補完する、疎性を意識したバッチング戦略を提案し、粗粒度の負荷分散を達成します。実験結果により、SparseBalance は LongBench ベンチマークにおいて、長文コンテキスト能力を 0.46
% 向上させつつ、エンドツーエンドで最大 1.33 imes の高速化を達成することが示されました。