1770146148
2026-02-03 19:05:00
研究者たちは、大規模言語モデル (LLM) トレーニングにおける重要な課題、つまりパフォーマンスを犠牲にすることなく再現性を維持することに取り組んでいます。上海交通大学およびByteDance SeedのXinwei Qiang氏、Hongmin Chen氏、Shixuan Sun氏は、Jingwen Leng氏、Xin Liu氏、Minyi Guo氏らとともに、一貫した結果を得るために不可欠な決定論的な注意が、一部の実装ではトレーニングのスループットを大幅に削減し、最大37.9%も削減できることを実証しました。彼らの研究では、決定論的アテンションをスケジューリング問題として定式化し、計算のボトルネックを最小限に抑えるスケジュールを提供する新しいアプローチである DASH (高スループットのための決定論的アテンション スケジューリング) が導入されています。操作の順序を最適化することで、DASH はスループットを最大 1.28 倍向上させることが実証されており、効率的で再現可能な LLM トレーニングに向けた大きな一歩となります。
チームは、LLM トレーニングにおける重要な課題、つまり、複数の実行にわたって決定論的でビット単位の同一の結果を保証することに関連する多大なパフォーマンス コストに対処しました。 FlashAttendant-3 などの広く使用されているアテンションの実装では、決定論的なバックワード パスを有効にすると、数値の一貫性のために必要な勾配累積演算のシリアル化により、スループットが最大 37.9% 削減される可能性があります。このパフォーマンスの低下は、計算フェーズと勾配低減フェーズの非効率なスケジューリングによって発生し、ハードウェア リソースの十分な活用につながります。
研究者らは、決定論的アテンション バックワード パスを有向非巡回グラフ (DAG) 上のスケジューリング問題として定式化し、クリティカル パスの長さを最小限に抑えるように設計されたスケジュールを導き出しました。この定式化に基づいて、彼らは、2 つの相補的なスケジューリング戦略を組み込んだフレームワークである DASH (高スループットのための決定的アテンション スケジューリング) を提示しています。 1 つ目の降順 Q タイル反復では、特に因果的アテンションにおけるパイプラインの停止を減らすために、逆クエリ ブロック トラバーサルを採用しています。 2 番目のシフト スケジューリングは、DAG モデル内で理論的に最適なスケジュールであり、フル アテンション マスクと因果的アテンション マスクの両方でパイプラインの停止を軽減します。
この革新的なアプローチは、タイルの実行を累積順序に合わせて調整するという中心的な問題、つまりパフォーマンス低下の主な原因として特定されている調整不良に取り組みます。 NVIDIA H800 GPU で実施された実証評価では、DASH が決定論的注意のパフォーマンス ギャップを効果的に狭めることが実証されています。提案された戦略により、アテンション バックワード パスのスループットがベースラインと比較して最大 1.28 倍向上し、再現可能な LLM トレーニングの効率が大幅に向上しました。この改善は、リダクション プロセスを並列化することで実現され、CTA が異なるタイルでリダクションを同時に開始できるようになり、連続操作によって引き起こされるボトルネックが回避されます。
この研究は、決定論的 LLM トレーニングの新しいベンチマークを確立し、より効率的で信頼性の高い大規模モデル開発への道を提供します。この調査では、パフォーマンスのギャップはシリアル化に固有のものではなく、最適ではないタイル スケジューリングと厳格な累積順序の結果であることが判明しました。決定論的なバックワード パスを DAG としてモデル化することで、チームはクリティカル パス長を最適化する戦略を設計することができ、よりバランスのとれたワークロードを確保し、シリアル リダクション操作中の競合を軽減できました。この DAG ベースの形式化は、スケジューリング プロセスの原則に基づいた最適化を可能にする新しい貢献を表しています。 DASH コードのオープンソース化により、LLM コミュニティ内でのさらなる研究と採用が促進されます。
DAG スケジューリングにより決定論的なアテンション バックワード パスを最適化
科学者たちは、再現性にとって重要な側面である大規模言語モデルのトレーニングにおける決定論的注意のパフォーマンスの限界に取り組みました。研究チームは、FlashAttendant-3 のようなアテンション実装における決定論的なバックワード パスでは、勾配累積のシリアル化により、非決定論的な対応物と比較してスループットが最大 37.9% 低下する可能性があることを確認しました。これを克服するために、彼らはクリティカル パス長を最小限に抑えることを目的として、有向非巡回グラフ (DAG) 上のスケジューリング問題として決定論的アテンション バックワード パスを定式化しました。この革新的なアプローチにより、2 つの相補的な戦略を組み込んだ DASH (高スループットのための決定的アテンション スケジューリング) の開発が可能になりました。
DASH の中核は、特に因果的注意メカニズムでパイプラインの停止を減らすように設計された逆クエリブロック トラバーサルである降順 Q タイル反復にあります。研究者は、クエリ ブロックを逆の順序で処理し、データ フローを最適化し、依存関係を減らすことでパイプラインの停止を縮小するためにこの方法を設計しました。これを補完するシフト スケジューリングは、DAG モデル内で理論的に最適なスケジュールとして開発され、フル アテンション マスクと因果的アテンション マスクの両方のパイプライン ストールを削減します。実験では、NVIDIA H800 GPU、CUDA 12.6、および Triton 3.4 を使用し、すべてのカーネルは FlashAttendant-3 実装を拡張することによって実装されました。
ベースライン比較は、FlashAttendant-3 の決定論的バックワード パスと Triton チュートリアルの因果的アテンションの実装に対して実行されました。この調査では、合計 16,384 個のトークンを固定し、シーケンス長を 512 から 16,384 まで変化させてパフォーマンスのベンチマークを行い、ヘッド ディメンション 64 と 128 で隠れディメンション 2,048 をテストしました。すべて BF16 高精度ランダム入力を使用しました。結果は、DASH によってアテンション後方パスのスループットがベースラインと比較して最大 1.28 向上し、再現可能な LLM トレーニングの効率が大幅に向上することが実証されました。詳細な分析により、シーケンス長 16,384、KV ブロック サイズ 128 で、計算が 128 個のストリーミング マルチプロセッサ (SM) に分散されていることが明らかになりました。
研究チームは、SM 間の通信遅延、特にリモート L2 キャッシュ セグメントへのアクセス (200 サイクルから 500 サイクル以上) が制限要因になることを観察しました。シフト スケジューリングの複雑な依存関係グラフは計算上の利点をもたらしましたが、極端な並列処理ではこの通信オーバーヘッドの影響をより受けやすく、場合によってはわずかなパフォーマンスの低下を示すことが判明しました。因果的アテンション マスクの場合、降順 Q タイル反復と対称シフト スケジューリングの両方で一貫してスループットが向上し、対称シフト スケジューリングはヘッド ディメンション 64 で優れたワークロード バランシングを実証しました。
DASH スケジューリングにより、確定的な LLM トレーニングのスループットが向上します
科学者たちは、決定論的大言語モデル (LLM) トレーニングにおけるパフォーマンスのボトルネックに対処するために、新しいスケジューリング フレームワークである DASH (高スループットのための決定論的アテンション スケジューリング) を開発しました。再現性にとって重要な決定論的トレーニングでは、多くの場合、大幅なパフォーマンス コストが発生し、FlashAttendant-3 などの広く使用されているアテンションの実装では、最大 37.9% のスループットの低下が観察されています。この損失は、数値の一貫性を確保するために勾配累積演算をシリアル化することで発生し、ハードウェアの使用率が低下します。研究者らは、クリティカル パスの長さを最小限に抑えるために、有向非巡回グラフ (DAG) 上のスケジューリング問題として決定論的アテンション バックワード パスを定式化しました。
実験の結果、DASH は決定論的注意のパフォーマンス ギャップを狭め、ベースラインと比較してスループットを最大 1.28 向上させることが明らかになりました。チームは NVIDIA H800 GPU でこの改善を測定し、再現可能な LLM トレーニングの効率が大幅に向上していることを実証しました。このブレークスルーにより、決定論的アテンション バックワード パスの大幅な高速化が実現し、数千の GPU にわたって LLM をスケーリングする際の重要な課題に対処できます。データによると、決定論的 FlashAttend のパフォーマンス低下は、タイルの実行と累積順序の間の不整合が原因で発生し、削減プロセスにボトルネックが生じていることが示されています。
この研究では、降順 Q タイル反復とシフト スケジューリングという 2 つの相補的なスケジューリング戦略が導入されています。逆クエリ ブロック トラバーサルである降順 Q タイル反復は、因果的アテンションにおけるパイプライン ストールを縮小します。一方、DAG モデル内の理論的に最適なアルゴリズムであるシフト スケジューリングは、フル マスクと因果的マスクの両方のパイプライン ストールを削減します。測定により、シフト スケジューリングが計算タスクの位相シフト割り当てを採用し、完全に時間差のある実行パターンを作成し、モデルの理論上の使用率限界に近づいていることが確認されています。テストでは、DASH が削減プロセスを効果的に並列化し、順次削減を強制する単純なスケジュールとは異なり、CTA が異なるタイルで同時に削減を開始できることが証明されています。研究では、決定論的注意におけるパフォーマンス低下の主な原因は、タイルの実行と累積順序間の不整合であることが特定されました。その結果、この研究は決定論的アテンション バックワード スケジューリングの DAG ベースの形式化を提供し、クリティカル パス長の原則に基づいた最適化を可能にし、最終的に LLM トレーニングの効率を向上させます。
DASH フレームワークは GPU での決定論的な注意を加速します
科学者たちは、大規模言語モデル (LLM) トレーニングの重要な側面である、注意メカニズムにおける決定論的な後方パスに関連するパフォーマンスのペナルティに取り組んできました。彼らは、有向非巡回グラフ (DAG) 上のスケジューリング問題として計算を定式化し、降順 Q タイル反復とシフト スケジューリングという 2 つのスケジューリング戦略を備えたフレームワークである DASH を導入しました。降順 Q タイル反復は、クエリブロックの逆トラバーサルを通じて因果関係の注意を加速します。一方、シフト スケジューリングは、DAG モデル内で理論的に最適なスケジュールを目指して、フル マスクと因果マスクの両方のパイプライン ストールを削減します。 H800 GPU での実証評価では、DASH が決定論的注意のパフォーマンス ギャップを大幅に狭め、ベースライン手法と比較してスループットが最大 1.28 倍向上することが実証されました。
この進歩は、より効率的で再現可能な LLM トレーニングに貢献します。著者らは、理論的な最適性が必ずしも実際的な優位性につながるわけではないことを認めており、レジスタープレッシャーや SM 間の通信遅延などのハードウェア制限がパフォーマンスに影響を与える重要な要素であると特定しています。今後の作業では、これらのハードウェアの現実を考慮して、さらなる最適化を検討する可能性があります。この調査では、帯域幅やメモリだけに焦点を当てるのではなく、実行と蓄積の順序を同時に最適化することの重要性が強調されています。 DASH は一連のソリューションを提供することで、実践者が LLM トレーニングの再現性を維持しながら、高いスループットの注意力を達成できるようにします。この調査結果は、この領域でパフォーマンスを最大化するには、理論的な最適性と実際のハードウェア制約のバランスをとる微妙なアプローチが不可欠であることを示唆しています。
#Dash #スケジューリングにより再現可能な #LLM #トレーニング #スループットが #高速化されました