※本記事は、Stanford CS336 Language Modeling from Scratch Spring 2025の講義動画「GPUs」の内容を基に作成されています。講義の詳細情報は https://stanford-cs336.github.io/spri... でご覧いただけます。Stanford大学のオンラインAIプログラムについては https://stanford.io/ai 、本講義への登録については https://online.stanford.edu/courses/c... をご参照ください。
本記事では、講義の内容を詳細にまとめておりますが、要約や解釈による誤りがある可能性もありますので、正確な情報や文脈については、オリジナルの講義動画をご視聴いただくことをお勧めいたします。
登壇者紹介
- Percy Liang - Stanford大学コンピューターサイエンス准教授、Center for Research on Foundation Models (CRFM) ディレクター
- Tatsunori Hashimoto - Stanford大学コンピューターサイエンス助教授
本講義はStanford Onlineを通じて提供されており、Stanford工学部のCenter for Global & Online Education (CGOE)によって運営・管理されています。Stanford Onlineは、Stanford大学全体の学部・部門が提供する学位プログラム、単位付与教育、専門証明書プログラム、無料公開コンテンツの包括的なカタログを提供し、世界中の学習者の知識拡充、キャリア向上、人生の充実を支援しています。
1. イントロダクション
1.1 授業の概要とFlash Attention 2の実装
皆さんがAssignment 1を楽しんでいることを願っています。今夜が締切です。延長が必要な場合はお知らせください。Assignment 2はもうすぐ公開予定で、現在Tritonの部分に最終的な仕上げを加えているところです。きっと楽しんでいただけると思います。Flash Attention 2またはFlash Attention 2の一部を実装していただくことになり、これは非常に興味深い内容になると考えています。
今日はGPUについて話します。GPUは私たちの言語モデルを動かすものであり、正しく理解することが極めて重要です。もしこれまでモデルを実行するハードウェアについて真剣に学んだことがない場合、GPUは非常に神秘的に見えるかもしれません。今日の私の目標は、CUDAとGPUをより理解しやすいものにすることです。
私が解明したいことの一つは、このスライドの図表を理解する必要はありませんが、なぜGPUが遅くなるのかということです。GPUは非常に神秘的な方法で遅くなります。この図表について講義の終わり近くで詳しく説明する予定です。行列乗算のサイズを増やしていくと、遅くなったり速くなったりするかもしれませんが、これらの非常に予測不可能に見える波のようなパターンが現れ、なぜGPUが特定の数の特定の倍数で速く、他の場合は遅いのか疑問に思うことでしょう。これは非常に神秘的です。これを理解しようと思います。
もう一つは、高速アルゴリズムの作り方を理解したいということです。ほぼ全員がFlash Attentionについて聞いたことがあると思います。これはTransformer内の注意操作を非常に巧妙に計算することで、はるかに長いコンテキストを可能にするものです。そして、おそらくFlash Attentionのような新しいアルゴリズムや新しい実装を思いつきたいと思うでしょう。そのためには、どのようなプリミティブとどのようなコンポーネントを理解する必要があるのでしょうか。
これらが今日の2つの学習目標です。1つ目は、講義の終わりまでにGPUについて快適に感じ、GPUがどのように動作するかを理解することです。2つ目は、アルゴリズムの特定の部分を高速化することに快適さを感じることです。新しいアーキテクチャを作る場合、CUDAでそれを高速化できると感じられるようになることを望んでいます。
ハードウェアは必ずしも私が働く分野ではないため、多くのクレジットを与える必要がある特別なリソースがあります。特にHorus HeathのブログがあR、そこで多くの楽しいGPUの事実を学ぶことができます。例えば、なぜゼロで満たされた行列乗算がそうでないものよりも速いのかということを、彼のブログに行くことで学ぶことができます。CUDAモードグループやGoogleの素晴らしいTPUブックなど、他にも私が引用したリソースがあります。このトピックに興味がある場合は、これらのリソースを見に行ってより多くのことを学ぶことをお勧めします。これはある意味では浅いですが、ハードウェアの完全なカバレッジを提供することを願っているからです。
1.2 GPUの謎の解明:性能の波状パターンの理解
私が解明したいことの一つは、このスライドの図表を理解する必要はありませんが、スライドには多くの情報があることは承知していますが、なぜGPUが遅くなるのかということです。GPUは非常に神秘的な方法で遅くなります。私はこの図表について講義の終わり近くで詳しく説明しようと思います。
行列乗算のサイズを増やしていくと、遅くなったり速くなったりするか、何らかの変化があるかもしれませんが、これらの非常に予測不可能に見える波のようなパターンが現れます。なぜ私のGPUが特定の数の特定の倍数で速く、他の場合は遅いのか疑問に思うことでしょう。これは非常に神秘的です。私たちはこれを理解しようと思います。
この謎の図表を完全に理解することが、今日の講義の重要な目標の一つです。講義の終わりまでに、皆さんはこの種の現象の各々を理解できるようになるはずです。皆さんは「ええ、その図表は完全に普通に見える。それはGPUが行う自然なことだ」と言えるようになるでしょう。
この波状パターンの謎を解くために、私たちはGPUがどのように動作し、なぜ特定の条件下で性能が劇的に変化するのかを深く理解する必要があります。この理解こそが、効率的なアルゴリズムを設計し、ハードウェアを最大限に活用するための基盤となるのです。
1.3 高速アルゴリズム開発の学習目標
もう一つは、高速アルゴリズムの作り方を理解したいということです。ほぼ全員がFlash Attentionについて聞いたことがあると思います。これはTransformer内の注意操作を非常に巧妙に計算することで、はるかに長いコンテキストを可能にするものです。
そして、おそらく皆さんはFlash Attentionのような新しいアルゴリズムや新しい実装を思いつきたいと思うでしょう。そのためには、どのようなプリミティブとどのようなコンポーネントを理解する必要があるのでしょうか。これらが今日の2つの学習目標の種類です。
1つ目は、講義の終わりまでにGPUについて快適に感じてほしいということです。GPUがどのように動作するかを理解してほしいのです。2つ目は、アルゴリズムの特定の部分を高速化することに快適さを感じてほしいということです。新しいアーキテクチャを作る場合、うまくいけばCUDAでそれを高速化しようと試みることができると感じられるようになってほしいのです。
ハードウェアは必ずしも私が働く分野ではないため、多くのクレジットを与える必要がある特別なリソースがあります。特にHorus HeathのブログがありR、そこで多くの楽しいGPUの事実を学ぶことができます。例えば、なぜゼロで満たされた行列乗算がそうでないものよりも速いのかということを、彼のブログに行くことで学ぶことができます。CUDAモードグループやGoogleの素晴らしいTPUブックなど、他にも私が引用したリソースがあります。このトピックに興味がある場合は、これらのリソースを見に行ってより多くのことを学ぶことをお勧めします。これはある意味では浅いですが、うまくいけば完全なハードウェアのカバレッジを提供できることを願っているからです。
2. コンピュート・スケーリングの背景
2.1 スケーリング法則と計算量の重要性
皆さんの多くはNLPコースを受講したことがあり、最近のNLPコースでは、ある程度のスケーリング法則を教えていると思いますので、おそらくこれを見たことがあるでしょう。これは単に文脈を設定するためのものです。
私たちは、より多くのcomputeを持つことが大規模言語モデルの訓練に役立つことを知っています。これは事前訓練のスケーリングチャートですが、お望みであれば推論スケーリングチャートに置き換えることもできます。より多くのcomputeを持てば持つほど、データに対してより多くの処理を行うことができるということが一般的に合意されています。
より多くのデータを取り込むことができ、より大きなモデルを訓練することができ、これらすべてが性能の向上につながります。そこで、もちろん深層学習は本当に重要だと考えるかもしれませんが、実際に性能を駆動しているのは、より高速なハードウェア、より良い利用率、改善された並列化です。これがハードウェアを理解することが重要である理由の舞台設定の一種です。
そしてもちろん、computeスケーリングについて考えると、私たちは「さて、どうやってcomputeスケーリングを得るのか?どうやってモデルをより速く訓練させるのか?」と問うことになります。
2.2 Dennard scalingからパラレル・スケーリングへの移行
半導体スケーリングの初期の頃を考えてみると、CPUがどのように高速化するかを考えていた場合、それらはDennard scalingと呼ばれるもので拡張していました。Moore's lawでは、チップ上のトランジスタの数を毎年2倍にしていました。この2倍化があると、より小さく、より小さなトランジスタがより速く、より速いクロック速度で、より低く、より低い電力で駆動できるようになり、それが順番により多くの性能をもたらすDennard scalingになります。
1980年代から2000年代にかけて、これは頭打ちになりました。HennessyとPattersonによるこのチャートで見ることができるように、単一スレッド性能、これがここの青い点ですが、基本的に先細りになり始めました。
もちろん、トランジスタの数は実際には減少し始めませんでした。より高く、より高いトランジスタ密度を持つチップがありましたが、それは役に立ちませんでした。単一スレッドでより高いスループットを提供していませんでした。
これは、絶対的な意味でより速く計算を行うことができないということを意味します。それを補うために必要なのは並列スケーリングです。深層学習とニューラルネットワークのスケーリングの物語は、単一スレッドスケーリング、つまり絶対的な意味でより速く計算を行うことから、多くのワークロードがすべて一度に計算される並列スケーリングへと移行することです。
2.3 GPU性能の超指数的増加(K20からH100まで)
これは私のお気に入りの、Bill Dowyの基調講演による、computeスケーリングチャートの一つです。彼は秒当たりの整数演算数の超指数的な増加を示しており、最初期のK20sからH100まで続いています。これは本当に驚くべき指数的または超指数的な曲線の一種です。
私たちは言語モデルから最大限を得るために、この曲線を活用する方法を本当に理解する必要があります。これが私たちの目標の一種になります。
この超指数的成長は、GPU技術の進歩がいかに急速で劇的であったかを示しています。K20から始まってH100に至るまでの軌跡は、単なる漸進的改善ではなく、真に変革的な性能向上を表現しています。この曲線を理解し活用することが、現代の言語モデル開発において決定的に重要なのです。
3. CPUとGPUの基本的な違い
3.1 設計思想の違い:レイテンシ最適化 vs スループット最適化
私はすでにこの重要な違いをほのめかしましたが、CPUは誰もがプログラミングを始めると馴染みのあるものだと思います。プログラムがあって、それが単一のスレッドで進行し、何が起こっているかをステップバイステップで実行するという実行モデルです。その種の実行モデルをサポートするために、何が必要でしょうか?大きな制御ユニットが必要です。多くの分岐と多くの条件制御ロジックがあるため、これらのものを非常に迅速に実行する必要があります。
CPUは、この抽象化された図で見ると、チップの多くを大きな制御、分岐予測に捧げ、それほど多くのスレッドを持たないため、これらを非常に迅速に実行します。現在は多くのコアを持つCPUがありますが、GPUと比較すると、ほとんど何もありません。
対照的に、GPUは本当にたくさんの計算ユニット、ALUを持っています。そこの小さな緑のボックスがそれらです。制御に捧げられるチップの量ははるかに少ないです。たくさんの計算ユニットを並列に動作させる、種類の制御ロジックを調整する少しの制御ロジックがあります。
これは、CPUとGPUで何が強調されているかの一種の図です。設計目標が何であるかを見ると、それらは非常に異なる種類の目標に向けて設計されています。CPUはレイテンシの最適化と考えることができます。私は自分のタスクをできるだけ早く終わらせたいのです。
右側にタスクT1からT4があるとすると、CPUでは、各タスクをできるだけ早く終わらせようとします。これらのタスクのうち任意の一つを早く終わらせたい場合、T1は本当に早く完了します。GPUでは、高いスループットを最適化しています。レイテンシは気にしません。集約的に持っているすべてのタスクをできるだけ早く完了させたいだけです。
それをサポートするために、多くのスレッドがあり、これらのスレッドは非常に迅速に眠りについたり目覚めたりすることができます。最終的に、ワークロードT1からT4をすべて、CPUのものより前に終了させることができますが、個別にはこれらすべてがより高いレイテンシを持っています。つまり、それらは異なる設計原則と設計目標を持っているのです。
3.2 ハードウェア構成の違い:制御ユニット vs 計算ユニット
CPUは、この抽象化された図で見ると、チップの多くを大きな制御、分岐予測に捧げ、それほど多くのスレッドを持たないため、これらを非常に迅速に実行します。現在は多くのコアを持つCPUがありますが、GPUと比較すると、ほとんど何もありません。
対照的に、GPUは本当にたくさんの計算ユニット、ALUを持っています。そこの小さな緑のボックスがそれらです。制御に捧げられるチップの量ははるかに少ないです。たくさんの計算ユニットを並列に動作させる、種類の制御ロジックを調整する少しの制御ロジックがあります。
これは、CPUとGPUで何が強調されているかの一種の図です。CPUでは、プログラムの逐次実行をサポートするために大きな制御ユニットが必要で、多くの分岐と多くの条件制御ロジックがあるため、これらのものを非常に迅速に実行する必要があります。一方、GPUは大量の計算ユニット(ALU)に重点を置き、それらを調整する制御ロジックは比較的小さく設計されています。
このハードウェア構成の違いが、それぞれの処理装置の特性と適用分野を決定する根本的な要因となっています。CPUは複雑な制御フローを効率的に処理できる一方、GPUは大量の並列計算を同時に実行することに特化しているのです。
3.3 実行モデルの違い:シーケンシャル vs パラレル
CPUは誰もがプログラミングを始めると馴染みのあるものだと思います。プログラムがあって、それが単一のスレッドで進行し、何が起こっているかをステップバイステップで実行するという実行モデルです。その種の実行モデルをサポートするために、何が必要でしょうか?大きな制御ユニットが必要です。多くの分岐と多くの条件制御ロジックがあるため、これらのものを非常に迅速に実行する必要があります。
対照的に、GPUの実行モデルは根本的に異なります。GPUは大量の計算ユニットを並列に動作させ、たくさんの計算ユニットを並列に動作させる種類の制御ロジックを調整する少しの制御ロジックがあります。
CPUは各タスクをできるだけ早く終わらせようとするシーケンシャルな処理に最適化されています。一方、GPUは高いスループットを最適化しており、レイテンシは気にしません。集約的に持っているすべてのタスクをできるだけ早く完了させたいだけです。それをサポートするために、多くのスレッドがあり、これらのスレッドは非常に迅速に眠りについたり目覚めたりすることができます。
最終的に、GPUではワークロードT1からT4をすべて、CPUのものより前に終了させることができますが、個別にはこれらすべてがより高いレイテンシを持っています。つまり、CPUとGPUは異なる設計原則と設計目標を持っているのです。この実行モデルの違いこそが、それぞれの処理装置が最適とする計算タスクの性質を決定しているのです。
4. GPU アーキテクチャ
4.1 ストリーミング・マルチプロセッサ(SM)とストリーミング・プロセッサ(SP)
GPUはかなり異なる解剖学的構造を持っています。皆さんが今までGPUのレイアウト図がどのようなものか見たことがあるかわかりませんが、実際にチップの図を少し後でお見せします。しかし、GPUの背後にある核心的なアイデアと重要な概念は、GPUが多くのSM、ストリーミング・マルチプロセッサを実行するということです。
ストリーミング・マルチプロセッサは、Tritonのようなもので プログラミングする際に、SMのレベルで動作することになるため、原子単位として考えることができます。各SM内には多くのSP、ストリーミング・プロセッサが含まれており、ストリーミング・プロセッサは多くのスレッドを並列に実行します。
これについて考える一つの方法は、SMは多くの制御ロジックを持っているということです。例えば、分岐を実行することを決定できます。SPは同じ命令を取り、それを多くの異なるデータに適用して動作します。このモデルの下で、大量の並列計算を行うことができます。
SMは制御の各粒度単位の一種です。SPは個別に多くの計算を行うことができます。A100を見ると、これは現時点で前世代のGPUですが、128のSMがあります。これはほとんどのCPUのコアよりもはるかに多くです。これらのSMのそれぞれには、非常に多数のSPと、それらの内部にある特殊化された行列乗算ユニットがあります。これが計算モデルの一種です。
ここで質問がありましたか?申し訳ありません。
学生からの質問がありました:「前のスライドのGPUは、このGPUと同じですか?」
はい、これはそれの一種の漫画版です。各行をSMとして考えることができます。それは独自の制御ユニットを持っています。各緑のブロックは、ここのこれらの緑のブロックの一つのようなSP32処理ユニットの一種かもしれません。各SMは、計算を行うために所有するテンソルコアのような様々な部品を動作させることができます。
4.2 A100の構成:128 SM、各SMに多数のSPと専用行列演算ユニット
A100を見ると、これは現時点で前世代のGPUですが、128のSMがあります。これはほとんどのCPUのコアよりもはるかに多くです。これらのSMのそれぞれには、非常に多数のSPと、それらの内部にある特殊化された行列乗算ユニットがあります。これが計算モデルの一種です。
A100の構成は、GPUの大規模並列処理能力を具現化したものです。128という数のSMは、従来のCPUが持つコア数と比較すると圧倒的に多く、これがGPUの並列処理における優位性の源泉となっています。各SMが独立して動作できることにより、GPUは同時に多数の計算タスクを処理することが可能になります。
さらに重要なのは、各SM内に配置された特殊化された行列乗算ユニットの存在です。これらのユニットは深層学習で頻繁に使用される行列演算を高速に実行するために設計されており、現代のニューラルネットワーク処理においてGPUが不可欠となっている理由の一つです。この特殊化されたハードウェアにより、GPUは汎用プロセッサでは実現困難な計算性能を達成しているのです。
4.3 物理的レイアウトとメモリ階層の重要性
計算と並んで、2つの重要なことを念頭に置く必要があります。GPUはコンピュータなので計算を行いますが、実際には計算は追跡する必要がある2つの重要なことのうちの1つに過ぎません。メモリは現時点では間違いなくより重要であり、GPUでプログラムを実行する際の性能プロファイルにおいて、より重要であり続けるでしょう。
メモリを理解するために、GPUとチップの物理的レイアウトを理解する必要があります。なぜなら、ある意味で、このような高速で動作している際には、メモリの物理的近接性がかなり重要になり始めるからです。そこで、物事がどのように配置されているかの物理的近接性と、それがメモリアクセスと性能についてどのように考えるべきかに関連する方法をお見せします。
各SMにより近いメモリは、より高速になります。L1と共有メモリのような非常に非常に高速な種類のメモリがあり、それはSM内に存在し、本当に高速です。レジスタのような物、頻繁に読み書きしている物は、L1と共有メモリに入れたいと思うでしょう。
ご覧のとおり、これらの緑の領域はSMです。そして、これらの青い領域があります。これはGPUチップ上にあり、これらはSMのすぐ隣にあるL2メモリです。SMの内部にはありませんが、物理的にはまだかなり近いです。これらはまだかなり高速です。10倍遅いですが、まだ合理的に高速です。
そして、チップ自体の外側に、これは3090カードか何かのようなもの、あるいはこれはPCI 100です。GPUがここにあり、実際にDRAMがチップの隣に存在しています。実際にチップの外に出て接続する必要があります。このチップ図でこれらの黄色いコネクタを端に見ることができます。これらはHPMコネクタです。これらは実際のGPUの外側にあるDRAMチップに接続しています。
これらにアクセスする速度を見ることができます。オンSMメモリは、そこから何かにアクセスするのに20クロックサイクルのようにはるかにはるかに高速ですが、L2キャッシュやグローバルメモリから何かにアクセスするのに200または300クロックサイクルのようなものがかかります。この10倍の要因があなたを本当にひどく傷つけるでしょう。
もしグローバルメモリにアクセスする必要がある計算があると、SMで行うべき作業が実際になくなってしまう可能性があります。すべての行列を乗算し終わって、今はただアイドル状態になる必要があります。そうすると利用率が良くなくなり、これがメモリについて考える本当に重要なテーマになるでしょう。メモリはある意味でGPUがどのように動作するかを考える鍵なのです。
5. GPUメモリ階層
5.1 物理的近接性と速度の関係
メモリを理解するために、GPUとチップの物理的レイアウトを理解する必要があります。なぜなら、ある意味で、このような高速で動作している際には、メモリの物理的近接性がかなり重要になり始めるからです。そこで、物事がどのように配置されているかの物理的近接性と、それがメモリアクセスと性能についてどのように考えるべきかに関連する方法をお見せします。
各SMにより近いメモリは、より高速になります。これは非常にシンプルな原則です。L1と共有メモリのような非常に非常に高速な種類のメモリがあり、それはSM内に存在し、本当に高速です。レジスタのような物、頻繁に読み書きしている物は、L1と共有メモリに入れたいと思うでしょう。
ご覧のとおり、これらの緑の領域はSMです。そして、これらの青い領域があります。これはGPUチップ上にあり、これらはSMのすぐ隣にあるL2メモリです。SMの内部にはありませんが、物理的にはまだかなり近いです。
そして、チップ自体の外側に、これは3090カードか何かのようなもの、あるいはこれはPCI 100です。GPUがここにあり、実際にDRAMがチップの隣に存在しています。実際にチップの外に出て接続する必要があります。このチップ図でこれらの黄色いコネクタを端に見ることができます。これらはHPMコネクタです。これらは実際のGPUの外側にあるDRAMチップに接続しています。
この物理的近接性の原則は、GPU設計における基本的な制約であり、高速処理を実現するための物理法則に基づいた必然的な結果なのです。信号が物理的距離を移動する時間が性能に直接影響するため、最も頻繁にアクセスされるデータは可能な限りSMに近い場所に配置される必要があります。
5.2 メモリタイプ別アクセス時間:L1/共有メモリ(20サイクル)vs L2/グローバルメモリ(200-300サイクル)
これらにアクセスする速度を見ることができます。オンSMメモリは、そこから何かにアクセスするのに20クロックサイクルのようにはるかにはるかに高速ですが、L2キャッシュやグローバルメモリから何かにアクセスするのに200または300クロックサイクルのようなものがかかります。
この10倍の要因があなたを本当にひどく傷つけるでしょう。もしグローバルメモリにアクセスする必要がある計算があると、SMで行うべき作業が実際になくなってしまう可能性があります。すべての行列を乗算し終わって、今はただアイドル状態になる必要があります。そうすると利用率が良くなくなります。
L2キャッシュは10倍遅いですが、まだ合理的に高速です。しかし、この速度差は深刻な性能への影響をもたらします。20クロックサイクル対200-300クロックサイクルという差は、単なる数値以上の意味を持ちます。これは実際の計算処理において、メモリアクセス待機がボトルネックとなり、高価な計算ユニットが無駄にアイドル状態になってしまうことを意味しているのです。
この圧倒的な速度差こそが、GPU プログラミングにおいてメモリ階層を意識した設計が不可欠である理由です。効率的なGPUプログラムは、可能な限りオンチップメモリを活用し、グローバルメモリアクセスを最小限に抑える必要があります。
5.3 利用率への影響:メモリ待機による処理ユニットのアイドル状態
もしグローバルメモリにアクセスする必要がある計算があると、SMで行うべき作業が実際になくなってしまう可能性があります。すべての行列を乗算し終わって、今はただアイドル状態になる必要があります。そうすると利用率が良くなくなり、これがメモリについて考える本当に重要なテーマになるでしょう。
このメモリ待機による利用率低下は、GPUの性能を根本的に制限する要因となります。高価な計算ユニットが、メモリからのデータ到着を待って何もせずに待機している状態は、ハードウェアリソースの深刻な無駄使いを意味します。SMは本来なら大量の並列計算を実行できる能力を持っているにも関わらず、200-300クロックサイクルものメモリアクセス待機時間により、その能力を発揮できません。
これはGPUの設計思想である高スループットを実現する最大の障害となります。多数の計算ユニットを並列動作させることでスループットを最大化しようとするGPUの基本戦略が、メモリアクセスの遅延により台無しになってしまうのです。メモリはある意味でGPUがどのように動作するかを考える鍵であり、効率的なGPU利用のためには、この利用率への影響を最小限に抑える戦略的なメモリ管理が不可欠なのです。
6. GPU実行モデル
6.1 3つの粒度:ブロック、ワープ、スレッド
Assignment 2では、GPUのための高性能コードを実際に書くことになるので、GPUが実際にどのように物事を実行するかの実行モデルについて実際に考える必要があります。これはやや複雑ですが、それほど非常に複雑ではありません。
考える必要がある3つの粒度があります。ブロック、ワープ、スレッドがあり、これは粒度が狭くなる順序です。ブロックはスレッドの大きなグループの一種であり、各ブロックはSMに割り当てられます。
これをSMが一種のワーカーだと考えてください。それは独自の自律的な単位であり、ブロックは処理のためにSMに割り当てられます。これが各粒度単位です。
次に、これらのブロック内には多くのスレッドがあります。各スレッドは実行する必要がある作業の一部の一種です。これらのスレッドが実行されるとき、それらはグループで実行されることになります。
これはワープと呼ばれるものです。ブロックを取り、それはスレッドのコレクションですが、そのブロックからスレッドを取り、それらは毎回32個の連続番号のスレッドのグループで実行されることになります。それが一種のワープと呼ばれるものです。
この図でここで何が起こっているかを見ることができます。多くのブロックがあります。各ブロックは異なるSMに割り当てられます。各ブロック内には多くの異なるワープがあります。各ワープは多くのスレッドで構成され、これらのスレッドはすべて異なるデータで同じ命令を実行します。これが実行モデルの一種です。
6.2 ブロックからSMへの割り当て
ブロックはスレッドの大きなグループの一種であり、各ブロックはSMに割り当てられます。これをSMが一種のワーカーだと考えてください。それは独自の自律的な単位であり、ブロックは処理のためにSMに割り当てられます。これが各粒度単位です。
この図でここで何が起こっているかを見ることができます。多くのブロックがあります。各ブロックは異なるSMに割り当てられます。
この割り当てメカニズムは、GPUの並列処理における作業分散の基本単位を表しています。各SMは独立して動作する自律的なユニットとして機能し、割り当てられたブロックを責任を持って処理します。これにより、GPUは多数のSMを並列に動作させることで、大規模な計算タスクを効率的に分散処理することが可能になります。
ブロックをSMへ割り当てるこの仕組みは、GPU全体のワークロード管理において中核的な役割を果たしており、各SMが最適に利用されるよう計算タスクを均等に分配することが性能向上の鍵となります。
6.3 ワープ:32個の連続番号スレッドによる同一命令実行
次に、これらのブロック内には多くのスレッドがあります。各スレッドは実行する必要がある作業の一部の一種です。これらのスレッドが実行されるとき、それらはグループで実行されることになります。
これはワープと呼ばれるものです。ブロックを取り、それはスレッドのコレクションですが、そのブロックからスレッドを取り、それらは毎回32個の連続番号のスレッドのグループで実行されることになります。それが一種のワープと呼ばれるものです。
各ブロック内には多くの異なるワープがあります。各ワープは多くのスレッドで構成され、これらのスレッドはすべて異なるデータで同じ命令を実行します。これが実行モデルの一種です。
このワープという概念は、GPUのSIMT(Single Instruction, Multiple Thread)実行モデルの核心を成しています。32個という固定サイズは、GPUハードウェアの物理的制約に基づいて決定されており、この32個のスレッドが完全に同期して同一の命令を実行することで、効率的な並列処理を実現しています。ワープ内のすべてのスレッドは歩調を合わせて動作するため、一つでも異なる処理パスを取るスレッドがあると、性能に深刻な影響を与える可能性があります。
これらのブロック、ワープ、スレッドが何であるかは、おそらく神秘的に思えるでしょうが、後でCUDAカーネルのようなものを設計する際の性能にとって重要な意味を持つでしょう。これを覚えておいてください。進むにつれて記憶を新たにします。
7. GPUメモリ・プログラミングモデル
7.1 論理的メモリ階層:レジスタ、ローカル、共有、グローバルメモリ
それはGPUの一種の論理的実行モデルでした。もしそれを理解すれば、GPUがどのように物事を実行するかを理解できます。GPUの論理的メモリモデルもあります。今、私は物理的ハードウェアを見せているのではありません。これはGPUのプログラミングについてどのように考えるかの一種です。
レジスタがあります。これらは本当に高速で、単一の数値タイプのストレージを保存するものです。ローカルメモリがあり、共有メモリがあり、グローバルメモリがあります。メモリ階層がより遅く、より遅く、より遅くなっていきます。あなたのコードはグローバルメモリに書き込むことができます。
あまり頻繁に使用されない定数メモリに書き込むこともできます。各スレッドは独自のレジスタと共有メモリにアクセスできます。しかし、ブロック間を横断する情報は、グローバルメモリに書き込む必要があります。これは実際にかなり重要です。
つまり、何かを実行するスレッドを書くときはいつでも、理想的には同じ少量のデータで動作しているということです。その少量のデータを共有メモリに読み込み、すべてのスレッドがその共有メモリにアクセスすることを非常に喜んでいます。それが終了し、完了します。それは素晴らしい実行モデルでしょう。
代わりに、あちこちのデータにアクセスする必要があるスレッドがある場合、それは非常に非常に遅いグローバルメモリにアクセスしなければならないでしょう。このテーマはGPUで動作するさまざまな方法について話すときに戻ってくるでしょう。
7.2 ブロック間通信のグローバルメモリ要件
各スレッドは独自のレジスタと共有メモリにアクセスできます。しかし、ブロック間を横断する情報は、グローバルメモリに書き込む必要があります。これは実際にかなり重要です。
このブロック間通信の制約は、GPU プログラミングにおける根本的な設計原則の一つです。各ブロックは独立したSMに割り当てられ、自律的に動作するため、異なるブロック間でデータを共有する唯一の方法はグローバルメモリを経由することです。これは、ブロック内での高速な共有メモリアクセスとは対照的に、200-300クロックサイクルという大幅に遅いアクセス時間を意味します。
この制約により、効率的なGPUプログラムを設計する際には、可能な限りブロック内で処理を完結させ、ブロック間通信を最小限に抑える戦略が重要になります。アルゴリズムの設計段階から、どのデータがブロック間で共有される必要があるかを慎重に検討し、グローバルメモリアクセスを戦略的に配置することが、性能最適化の鍵となるのです。
この要件は、GPU の並列処理アーキテクチャの物理的制約から生じる必然的な結果であり、プログラマーはこの制約を理解した上で効率的なアルゴリズムを設計する必要があります。
7.3 共有メモリ活用の重要性
つまり、何かを実行するスレッドを書くときはいつでも、理想的には同じ少量のデータで動作しているということです。その少量のデータを共有メモリに読み込み、すべてのスレッドがその共有メモリにアクセスすることを非常に喜んでいます。それが終了し、完了します。それは素晴らしい実行モデルでしょう。
代わりに、あちこちのデータにアクセスする必要があるスレッドがある場合、それは非常に非常に遅いグローバルメモリにアクセスしなければならないでしょう。このテーマはGPUで動作するさまざまな方法について話すときに戻ってくるでしょう。
この理想的な実行モデルと問題のある実行モデルの対比は、GPU プログラミングの成功と失敗を分ける重要な要因です。共有メモリを効果的に活用できれば、全てのスレッドが20クロックサイクルという高速アクセスの恩恵を受けることができます。一方、データアクセスパターンが最適化されていない場合、頻繁なグローバルメモリアクセスにより200-300クロックサイクルの遅延が発生し、せっかくの並列処理能力が台無しになってしまいます。
共有メモリの活用は単なる最適化手法ではなく、GPUの根本的なアーキテクチャを理解し、それに適合したプログラム設計を行うための必須要件なのです。この概念は、後に説明するタイリングやその他の最適化技術の基盤となる重要な理解です。
8. TPUとの比較
8.1 テンソルコア、スカラーユニット、ベクターユニット、MXU構成
ここでちょっとした脇道です。昨年はTPUに関するリソースが少し不足していたため、これをカバーしませんでした。しかし、講義の最初に言及した素晴らしいTPUブックやインターネットウェブサイトが出てきました。それには実際に多くの素晴らしい詳細があり、私はTPUについて何人かのGoogleの人々と話しました。高レベルでは、それはGPUと非常に非常に似ており、TPUについて少し話したいと思います。皆さんはTPUで動作することは決してないかもしれませんが、これらの代替アクセラレータが多くの方法で非常に似ているように動作することを理解することが重要だと思います。
ここにTPUがどのように見えるかの図があります。テンソルコアと呼ばれるものがあり、精神的にはテンソルコアをSMまたはストリーミング・マルチプロセッサに似ていると考えることができます。これらのそれぞれは、データで動作できる独自の原子単位の一種です。
基本的に制御ユニットであるスカラーユニットがあり、CPU のような任意のことも行うことができます。ベクターがあり、それに対してエントリ別に動作したい場合、それを行うのに良い場所であるベクターユニットがあります。そして、MXUと呼ばれる行列乗算を行うことだけに専念したチップの非常に大きな特殊化された部分があります。
それから、ベクターメモリとSMEの両方である非常に高速なメモリがあります。これらの両方が、テンソルコア上の非常に高速なオンチップまたはテンソルコアメモリです。そして、チップの外に存在する高帯域幅メモリがあります。
SMとの類似性が見えることを願っています。外部に遅いメモリがあり、内部に非常に高速なメモリがあり、行列乗算を行う特殊化されたハードウェアがあります。コア構造は非常に同じです。
8.2 GPUとの共通点:高速オンチップメモリと低速オフチップメモリ
SMとの類似性が見えることを願っています。外部に遅いメモリがあり、内部に非常に高速なメモリがあり、行列乗算を行う特殊化されたハードウェアがあります。コア構造は非常に同じです。
TPU においても、ベクターメモリとSMEの両方である非常に高速なメモリがあります。これらの両方が、テンソルコア上の非常に高速なオンチップまたはテンソルコアメモリです。そして、チップの外に存在する高帯域幅メモリがあります。
この構造的類似性は偶然ではありません。高性能計算において、メモリ階層の設計は物理法則と性能要件に基づく必然的な帰結なのです。最も頻繁にアクセスされるデータは物理的に計算ユニットに近い場所に配置し、大容量だが低速なメモリは外部に配置するという基本原則は、GPUとTPUの両方で共通しています。
この共通のアーキテクチャパターンにより、GPU向けに開発された最適化技術の多くが、TPUのような他のアクセラレータにも適用可能となります。メモリアクセスパターンの最適化、データの局所性の活用、計算とメモリアクセスのバランスといった根本的な設計原則は、プラットフォームを超えて有効なのです。
8.3 特化設計:行列乗算専用最適化
違いは、来週の並列性に関する講義で話しますが、アクセラレータがどのように一緒になっているかが少し異なることです。そして、私はワープについて話しませんでした。そのような他のことについては話しませんでした。
テンソルコアは、行列乗算を行うことだけに最適化されているため、ある意味で非常にシンプルです。GPUとは異なり、テンソルコアはそれ以外のことを行おうとしません。そして、それはある意味で非常に非常にシンプルです。アーキテクチャははるかにシンプルですが、概念的には同じことを行っています。
学生からの質問がありました:「それがテンソルと呼ばれるのは、任意のテンソルで動作するように最適化されているからですか、それとも単に動作するだけですか?」
それは任意のテンソルで動作できますが、MXUが実行する操作は行列乗算であり、したがって常にテンソルで動作するバッチ行列乗算のようなものになります。つまり、それは両方へのイエスとノーの答えの一種です。それはテンソルで動作しますが、常に実行する操作は、あなたが行うことができるより複雑なテンソル操作ではなく、行列乗算です。
この特化設計の哲学は、TPUの最大の強みの一つです。汎用性を犠牲にすることで、行列乗算という特定の操作において極めて高い性能を実現しています。深層学習のワークロードの大部分が行列乗算で構成されていることを考えると、この設計判断は非常に合理的です。GPUが様々な計算タスクに対応しようとするのに対し、TPUは意図的に範囲を限定することで、その分野での最適化を徹底的に追求しているのです。
9. GPU成功の要因
9.1 スケーラビリティ:SMの追加による処理能力向上
GPUが非常に成功している理由は、本当に簡単にスケールアップできることです。より多くの処理能力が欲しければ、SMをもっと追加するだけです。クロックをより速く駆動したり、より多くの熱放散問題に対処したりすることを心配する必要がありません。
このスケーラビリティの優位性は、GPU設計の根本的な強みです。従来のプロセッサがクロック速度の向上で性能を向上させようとすると、発熱と電力消費が指数的に増加し、物理的な限界に直面します。しかし、GPUは水平スケーリングのアプローチを採用することで、この問題を回避しています。
SM を追加するアプローチは、各SMが独立して動作するGPUのアーキテクチャ特性により可能になります。新しいSMを追加しても、既存のSMの動作に影響を与えることがなく、線形的な性能向上を実現できます。この設計により、GPU メーカーはより多くのSMを搭載した新世代のGPUを比較的容易に開発でき、継続的な性能向上を実現しているのです。
この水平スケーリングの能力こそが、GPUが深層学習の時代において中心的な役割を果たしている理由の一つであり、計算需要の急激な増加に対応できる唯一の現実的なソリューションとなっているのです。
9.2 プログラマビリティ:CUDAの相対的な親しみやすさ
プログラミングの観点から、CUDAは威圧的ですが、実際にはそのプログラミングモデルのおかげで、それほどひどくプログラムすることではありません。それが動作する方法は、各SM内で、スレッドがあり、それが多くの異なるデータで同じ命令を実行するということです。これは概念的に理解するのが簡単です。それが何を意味するかを考えることができ、特に行列上で動作していて、非常にシンプルな操作を行っている場合、それはまさにこの種のSIMTモデルです。
CUDAの親しみやすさは、その抽象化レベルの適切なバランスにあります。低レベル過ぎて理解不可能になることもなく、高レベル過ぎて制御が効かなくなることもありません。各SM内でのスレッド実行モデルは、プログラマーにとって直感的に理解できる概念です。特に、同一命令を異なるデータに適用するというSIMTパラダイムは、行列演算のような数値計算において自然に適合します。
このプログラミングモデルの理解しやすさが、GPU コンピューティングの普及を大きく後押ししました。研究者や開発者が比較的短期間でGPU プログラミングの基本を習得でき、高性能計算の恩恵を享受できるようになったのです。この親しみやすさこそが、GPUが単なるグラフィックス処理装置から汎用並列計算プラットフォームへと進化する上で決定的な役割を果たしたのです。
9.3 軽量スレッド:高い利用率実現のための停止・開始機能
最後に、これらのスレッドのそれぞれは非常に軽量であり、いつでも停止して開始することができます。そのため、別のスレッドを待つ必要がある場合や、何かを退避させて別のプロセスを開始する必要がある場合、これらのスレッドはすべて非常に軽量です。これは、スレッドに関連する状態があまりないことを意味し、停止して開始することができ、これによりGPUは各SM内で高い利用率を得ることができます。
この軽量スレッドの特性は、GPU の高い利用率を実現する上で極めて重要な機能です。従来のCPU スレッドでは、コンテキストスイッチに大きなオーバーヘッドが伴いますが、GPU スレッドは最小限の状態しか保持しないため、ほぼ瞬時に切り替えることができます。
この能力により、あるスレッドがメモリアクセス待機で停止している間に、他のスレッドが即座に実行を開始できます。200-300クロックサイクルのグローバルメモリアクセス待機時間も、この軽量スレッドシステムにより効果的に隠蔽されます。待機中のスレッドを一時停止し、準備完了の他のスレッドを実行することで、計算ユニットが常に有用な作業を行い続けることができるのです。
この柔軟なスケジューリング能力こそが、GPUが数千から数万のスレッドを効率的に管理し、高いスループットを維持できる理由です。スレッドの軽量性と高速な切り替え能力により、各SMの利用率を最大化し、GPUの並列処理能力を最大限に活用することが可能になっているのです。
10. 行列乗算の特別な地位
10.1 初期の研究:グラフィックスハードウェアでの高速行列乗算
GPUは、明らかにグラフィックス処理ユニットです。その生涯の多くの間、初期の頃は、科学計算を行うために使用されていませんでした。しかし、それがプログラマブルだったため、研究者たちは早期のNVIDIA GPUを使用して高速行列乗算を行う方法を見つけ出しました。
これは、グラフィックスハードウェアで高速行列乗算を行うことに関する初期の論文の一つです。テクスチャバッファなどをハッキングして行列乗算を行う方法を示しており、行列乗算への特別なサポートなしでも、研究者たちはそれを行う方法を見つけ出しました。
この初期の研究は、GPU コンピューティングの歴史において画期的な意味を持ちます。本来グラフィックス処理のために設計されたハードウェアを、巧妙なプログラミング技術により数値計算に転用するという発想は、その後のGPU 発展の方向性を決定づけました。研究者たちは、テクスチャメモリやピクセルシェーダーといったグラフィックス専用機能を、行列データの保存と計算に流用する創意工夫に富んだ手法を開発しました。
この「ハッキング」的なアプローチは、GPU が本質的に並列計算に適したアーキテクチャを持っていることを実証しました。グラフィックス処理における大量のピクセル並列処理と、行列乗算における要素別並列計算の間に根本的な類似性があることが明らかになったのです。この洞察が、後にNVIDIA が CUDA のような汎用GPU コンピューティングプラットフォームを開発する動機となったのです。
10.2 V100以降のテンソルコア導入による性能ギャップ
しかし今では、特にこの時代では、NVIDIAや他の企業が行列乗算は特別だと認識しています。深層学習を行っている場合、ワークロードのほとんどは行列乗算であり、行列乗算はある意味で祝福された操作です。
これは、NVIDIA GPU の異なる世代による秒当たりテラフロップス数を示すチャートです。オレンジの線は、行列乗算を行っている場合に得られる性能です。青い線は、非行列乗算FLOPSです。V100でテンソルコアと呼ばれる行列乗算を行う特殊化されたハードウェアを入れ始めた時に、この大きな大きなギャップを見ることができます。
行列乗算性能と非行列乗算性能に関して、この巨大なギャップを見ることができます。そのため、どのような種類のニューラルアーキテクチャを設計する場合でも、アーキテクチャの部分でも言っていましたが、ワークロードのほとんどが行列乗算である必要があります。なぜなら、それはGPUで行うことができる他の操作よりも桁違いに速いものだからです。非行列乗算ベースのニューラルネットワークを作ると、大きな大きな問題に陥ることになります。
このV100以降のテンソルコア導入は、深層学習の発展において転換点となりました。専用ハードウェアによる行列乗算の劇的な高速化により、行列乗算とその他の演算の間に桁違いの性能差が生まれたのです。オレンジ線と青線の間の巨大なギャップは、単なる数値的な違いではなく、アルゴリズム設計における根本的な制約となります。
この性能ギャップは、ニューラルネットワークアーキテクチャの設計思想を根本的に変えました。効率的なモデルを作るためには、可能な限り多くの計算を行列乗算として表現する必要があります。畳み込み、全結合層、注意機構など、現代の深層学習における主要な構成要素が全て行列乗算を中核とするのは、この性能特性によるものです。逆に、要素別演算や複雑な制御フローに依存するアーキテクチャは、このハードウェア特性により深刻な性能ペナルティを受けることになるのです。
10.3 実験結果:行列乗算 vs 非行列乗算のTFLOPS比較
これは、NVIDIA GPU の異なる世代による秒当たりテラフロップス数を示すチャートです。オレンジの線は、行列乗算を行っている場合に得られる性能です。青い線は、非行列乗算FLOPSです。V100でテンソルコアと呼ばれる行列乗算を行う特殊化されたハードウェアを入れ始めた時に、この大きな大きなギャップを見ることができます。
行列乗算性能と非行列乗算性能に関して、この巨大なギャップを見ることができます。V100以前では、オレンジ線と青線はほぼ同様の性能を示していました。しかし、テンソルコアの導入により、行列乗算性能は劇的に向上し、非行列乗算性能との間に前例のない差が生まれました。
この実験結果が示すTFLOPS比較は、現代のGPU設計における行列乗算の特権的地位を明確に実証しています。数値的には、行列乗算専用ハードウェアにより、同じチップ上で桁違いの性能差が発生しているのです。この差は理論値ではなく、実際の計測結果に基づく実証的なデータです。
このチャートが深層学習コミュニティに与えた影響は計り知れません。研究者やエンジニアは、この性能特性を理解し、アルゴリズムとアーキテクチャの設計において行列乗算を最大限活用する必要性を認識するようになりました。効率的なモデルの設計は、もはや計算複雑度の理論的分析だけでなく、このような実際のハードウェア特性を考慮した実用的な観点が不可欠となったのです。
11. GPUコンポーネントのスケーリング比較
11.1 接続性(青線):PCIe、NVLinkの緩やかな成長
最後に知っておいてほしい一般的な事実として、行列乗算が速いということが一つありますが、もう一つ覚えておくことが重要なのは、GPUの異なるコンポーネントの相対的なスケーリングの種類です。
これは、LM訓練スタックと呼ぼうか、GPUの異なるコンポーネント、または異なるコンポーネントがどれだけ速くスケーリングしているかを示す非常に素晴らしいチャートです。
青い線は、GPUからホストへの接続性です。つまり、それが接続されているサーバーです。PCIe、NVLink、これらすべての高級相互接続を使用することができます。それらは成長していますが、やや遅く成長しています。このチャートは正規化されたスケーリング、つまり第1世代の相互接続と比較した帯域幅です。
青い線で示される接続性の成長は、他のコンポーネントと比較して著しく制限されています。PCIe、NVLink などの相互接続技術は継続的に改善されているものの、その成長率は相対的に緩やかです。これは、物理的な信号伝送の制約、電力効率の要件、および互換性の維持といった工学的制約によるものです。
この緩やかな成長は、システム全体の性能において重要な意味を持ちます。GPU 自体の計算能力が急激に向上しても、ホストシステムとの間のデータ転送が制約となり、全体的な性能向上が制限される可能性があります。特に、大量のデータをCPU メモリとGPU メモリ間で転送する必要があるアプリケーションでは、この接続性の制約がボトルネックとなることがあります。このため、効率的なGPU アプリケーション設計では、ホストとデバイス間のデータ転送を最小限に抑える戦略が重要になるのです。
11.2 グローバルメモリ速度(緑線):GDDR→HBM2Eで100倍向上も依然緩慢
緑の線、これはグローバルメモリ速度です。GDDRからHBM2Eに移行し、それははるかにはるかに速いです。これは対数スケールで100倍速いですが、これは依然として遅いスケーリングの一種です。
グローバルメモリ速度の100倍向上は、絶対的な数値としては印象的に見えますが、対数スケールで表示されているこのチャートでは、他のコンポーネントの成長と比較すると相対的に制限されていることが分かります。GDDRからHBM(High Bandwidth Memory)、そしてHBM2Eへの進化は、メモリ技術における重要な進歩を表していますが、それでもなお「遅いスケーリング」と分類されるのです。
この制限されたスケーリングには、物理的および技術的な理由があります。DRAM技術は基本的に電荷の蓄積と読み出しに基づいており、物理法則によって制約されています。メモリセルの微細化、アクセス速度の向上、消費電力の削減といった複数の要求を同時に満たすことは、技術的に極めて困難です。また、大容量メモリを実現するためには、アクセス速度との間でトレードオフが生じます。
この相対的に緩やかなメモリ速度の向上が、後に説明するメモリボトルネックの根本原因となっています。計算能力の爆発的成長に対してメモリ帯域幅の成長が追いつかないため、効率的なGPU プログラミングではメモリアクセスパターンの最適化が決定的に重要になるのです。
11.3 計算性能(灰線):1〜100,000倍の驚異的成長
灰色の線がここにあります。これは計算スケーリングです。これは、行列乗算FLOPSを考慮している場合の浮動小数点演算数です。これがどれだけ速くスケーリングしているかは驚異的です。1から100,000倍速いようなものです。
この灰色の線が示す計算性能の成長は、まさに驚異的という言葉がふさわしい軌跡を描いています。1から100,000倍という数値は、単なる改善ではなく、根本的な変革を表しています。これは対数スケールでの表示にもかかわらず、他のすべてのコンポーネントを大幅に上回る急激な上昇を示しています。
この驚異的な成長は、複数の技術的要因の組み合わせによって実現されています。プロセス技術の微細化、アーキテクチャの最適化、そして最も重要なのは行列乗算専用ハードウェア(テンソルコア)の導入です。特に、V100以降のテンソルコア世代では、専用回路による行列乗算の高速化により、計算性能は文字通り桁違いの向上を遂げました。
この100,000倍という数字は、わずか10年程度の間に達成された成果であり、コンピューティング史上でも稀に見る急激な性能向上です。Moore's lawによる従来の半導体進歩とは質的に異なる、アーキテクチャレベルでの革新が生み出した結果なのです。この計算能力の爆発的成長こそが、現代の深層学習における大規模モデルの実現を可能にした根本的な要因です。
11.4 ボトルネックの変遷:FLOPSからメモリへ
初期のスケーリングでは、問題はFLOPSベースであったかもしれません。行列乗算を行うのに十分なFLOPSがなかったのです。しかし今では、右端のH100まで来ると、これらは驚異的に速いGPUですが、ボトルネックはおそらくメモリになるでしょう。メモリがそれほど速く成長していないからです。
そして将来に向かうにつれて、これは実際には変わることはありません。DRAMはスケールするのが非常に困難です。このより大きく、より大きなギャップを持ち続けることになります。そのため、ハードウェア効率的なアルゴリズムを設計している場合は、メモリについてより多く、より多く考える必要があります。
このボトルネックの変遷は、GPU コンピューティングにおける根本的なパラダイムシフトを表しています。初期の時代では、十分な計算能力がないことが制約でした。研究者やエンジニアは、より多くのFLOPSを確保することに注力していました。しかし、テンソルコアの導入により計算能力が爆発的に向上した結果、制約の性質が完全に変化したのです。
現在のH100のような最新GPUでは、計算ユニットは膨大な処理能力を持っていますが、それを支えるメモリ帯域幅の成長が追いついていません。結果として、多くの実際のワークロードでは、計算ユニットがメモリからのデータ到着を待ってアイドル状態になってしまいます。これは、高価で高性能な計算リソースが十分に活用されない状況を意味します。
将来への展望も同様に重要です。DRAM技術の物理的制約により、メモリ帯域幅の改善は今後も制限され続けるでしょう。一方、計算能力は引き続き急激に向上すると予想されるため、このギャップはさらに拡大します。これは、効率的なアルゴリズム設計において、メモリアクセスパターンの最適化が従来以上に決定的に重要になることを意味しているのです。
12. GPUパフォーマンス最適化の謎解き
12.1 謎の波状パフォーマンスパターンの分析目標
さて、皆さんは全員GPU専門家になったので、機械学習ワークロードをGPU上で非常に高速に動作させたいと思います。
そこで、このチャートから始めて、私たちの目標の一つは、このチャートが正確に何であるかを理解することです。これは私たちのやる気を起こさせる良いパズルになると思います。
ここで行っているのは、正方行列を一緒に乗算することです。x軸は私の正方行列乗算のサイズです。y軸は、私が行っている1秒当たりの演算数です。これをy軸のハードウェア利用率として考えることができます。
より大きく、より大きな行列を取得するにつれて、より良く、より良いハードウェア利用率を得ることになります。なぜなら、行うべき作業がより多くあり、ジョブの起動などのオーバーヘッドに圧倒されないからです。
しかし、起こっているこれらすべての奇妙なことがあります。1、2、3つの異なる、4つの異なる線が見えます。これらの線のそれぞれは、非常に予測不可能に見える方法で波状になっています。そこで、これらの線で正確に何が起こっているのかを理解したいと思います。
このセクションの終わりまでに、私の約束は、皆さんがこれらの現象のそれぞれを理解できるようになることです。「ええ、その図表は完全に普通に見える。それはGPUが行う自然なことだ」と言えるようになるでしょう。
この謎の波状パターンは、GPU性能の複雑さを象徴的に表現しています。一見すると不規則で予測不可能に見えるこれらの波は、実際にはGPUアーキテクチャの様々な特性が相互作用した結果なのです。行列サイズの増加に対して線形的な性能向上を期待するのが自然ですが、実際には4つの異なる性能曲線が示すように、現実ははるかに複雑です。
12.2 Rooflineモデル:メモリ制限領域とスループット制限領域
最初の部分は、その図表を見ると、このような感じに少し見えることに気づくでしょう。システムハードウェアコースを受講したことがある場合は、これをrooflineモデルの一種として覚えているはずです。
rooflineモデルは基本的に、スループットまたは利用率を見ている場合、2つのレジームがあることを言っています。左側のこの曲線の緑の部分にある、メモリ制限されたレジームがあります。そして、右側にスループット制限された部分があります。
ある意味で、右側では計算ユニットを完全に利用していると考えることができます。すべての行列乗算ユニットが常に乗算しています。そして、ここの対角線では、何らかのメモリボトルネックがあり、持っている1バイト当たりのFLOPS量、算術強度によって計算を行う能力が制限されています。
私たちは、メモリに制約されているこの左側の領域にいることを避け、すべての計算ユニットをある意味で完全に利用している右側にいたいと思います。それが目標であり、うまくいけばこのrooflineモデルはこのようなものに見えるでしょう。この対角線の部分と、ここの一番上の平らな部分があります。
これがミステリーの一部です。このrooflineモデルは、GPU性能理解における基本的なフレームワークを提供します。メモリ制限領域では、データをメモリから計算ユニットに供給する速度が性能を決定します。いくら強力な計算ユニットがあっても、データが到着しなければ意味がありません。一方、スループット制限領域では、メモリ帯域幅は十分であり、計算ユニットの能力が性能を決定します。
この2つの領域の境界は、アルゴリズムの算術強度(演算数対メモリアクセス数の比)によって決まります。効率的なGPUプログラムの設計では、可能な限りスループット制限領域で動作するよう、算術強度を高める必要があります。
12.3 4つの異なるパフォーマンス曲線の存在
しかし、起こっているこれらすべての奇妙なことがあります。1、2、3つの異なる、4つの異なる線が見えます。これらの線のそれぞれは、非常に予測不可能に見える方法で波状になっています。そこで、これらの線で正確に何が起こっているのかを理解したいと思います。
このセクションの終わりまでに、私の約束は、皆さんがこれらの現象のそれぞれを理解できるようになることです。「ええ、その図表は完全に普通に見える。それはGPUが行う自然なことだ」と言えるようになるでしょう。
4つの異なるパフォーマンス曲線の存在は、GPU性能の複雑さを如実に示しています。単純なrooflineモデルであれば、滑らかな対角線から平坦な最大性能へと移行する単一の曲線を期待するところですが、実際には複数の層に分かれた波状パターンが現れています。
これらの異なる線は、行列サイズがGPUの様々なハードウェア特性とどのように相互作用するかを反映しています。各線は異なる最適化レベルや制約条件を表しており、同じ計算タスクでも行列の次元によって全く異なる性能特性を示すことを意味します。
この予測不可能に見える波状パターンこそが、GPU最適化を困難にする要因の一つです。単純に行列サイズを大きくすれば性能が向上するわけではなく、特定のサイズで急激な性能低下や向上が発生します。これらの現象を理解することで、効率的なGPUプログラムを設計するための実践的な知識を得ることができるのです。
13. 最適化技術1:条件分岐の回避
13.1 SIMT実行モデルでの条件分岐の問題
最初にお話ししたいのは条件分岐についてです。前述したように、GPUの実行モデルはSIMT、つまりSingle Instruction Multi-Threadと呼ばれるものです。ワープ内のすべてのスレッドは同じ命令を実行し、それを異なるデータに対して行います。
そこで、このようなコードを書いたとしたらどうなるでしょうか?if文があり、スレッドインデックスが4未満なら何かを行う、スレッドインデックスが4以上なら別のことを行うという、非常に単純な条件モデルがあるとします。
これをGPUで実行すると何が起こるでしょうか?4つのスレッドでA命令を実行することになります。else部分を実行することになっている他の4つのスレッドは実際に一時停止します。そして、これらの他の4つのスレッドが生き返ってXを実行し、最初の4つのスレッドが眠りにつき、これらの命令を交互に実行することになります。
なぜでしょうか?私はAとXを同時に異なるスレッドで実行できません。再び言いますが、すべてのスレッドは同じ命令を実行しなければなりません。
SIMT実行モデルでの条件分岐の問題は、GPUアーキテクチャの根本的な制約から生じています。32個のスレッドからなるワープは、物理的に同一の命令デコーダーと制御ユニットを共有しているため、異なる命令を同時に実行することは不可能です。条件分岐により実行パスが分岐すると、ハードウェアは一方のパスを実行している間、他方のスレッドを強制的に待機させる必要があります。
この制約は、CPUの分岐予測とは根本的に異なります。CPUでは、分岐予測により効率的に条件分岐を処理できますが、GPUでは物理的な制約により、分岐したパスを逐次実行せざるを得ません。結果として、並列性の恩恵を完全に失い、実質的にシーケンシャルな実行になってしまうのです。
13.2 ワープ内分岐によるスレッド休眠の実例
これをGPUで実行すると何が起こるでしょうか?4つのスレッドでA命令を実行することになります。else部分を実行することになっている他の4つのスレッドは実際に一時停止します。
そして、これらの他の4つのスレッドが生き返ってXを実行し、最初の4つのスレッドが眠りにつき、これらの命令を交互に実行することになります。なぜでしょうか?私はAとXを同時に異なるスレッドで実行できません。再び言いますが、すべてのスレッドは同じ命令を実行しなければなりません。
この具体例は、SIMT実行モデルの制約を明確に示しています。スレッドインデックスが4未満の4つのスレッドがA命令を実行する際、残りの4つのスレッドは物理的に同じ実行ユニットを共有しているため、何もせずに待機するしかありません。ハードウェアは異なる命令を同時に発行できないため、一方のグループが完了するまで他方は完全に休眠状態になります。
次に、スレッドインデックスが4以上の4つのスレッドがX命令を実行する番になると、今度は最初の4つのスレッドが休眠します。この交互実行パターンにより、本来8つのスレッドで並列実行できるはずの処理が、実質的に4つのスレッドずつのシーケンシャル実行になってしまいます。
この現象は「warp divergence」と呼ばれ、GPU性能を大幅に低下させる主要因の一つです。32個のスレッドを持つワープでより複雑な分岐が発生すると、さらに多くのスレッドが無駄に待機することになり、並列処理の利点が完全に失われてしまうのです。
13.3 パフォーマンス低下メカニズム
なぜでしょうか?私はAとXを同時に異なるスレッドで実行できません。再び言いますが、すべてのスレッドは同じ命令を実行しなければなりません。
そのため、単一のワープ内での条件文は本当に本当に損害を与える可能性があります。なぜなら、まさにメインの制御フロー実行を行っていないスレッドを一時停止することを強制するからです。
それが、私が言及したかった唯一の非メモリ関連のことでした。大規模並列計算ユニットに条件文を入れるべきではないということは、かなり明白であるはずです。
しかし、それを片付けたら、考慮する必要がある他のトリックは、すべてメモリベースの一種です。
このパフォーマンス低下メカニズムは、GPUの物理的アーキテクチャに起因する避けられない現象です。ワープ内の32個のスレッドは同一の命令デコーダーと実行ユニットを共有しているため、異なる実行パスに分岐すると、ハードウェアは物理的に両方を同時に処理できません。結果として、並列実行の利点が完全に失われ、最悪の場合はシーケンシャル実行と同等の性能まで低下してしまいます。
この問題の深刻さは、分岐の複雑さに比例して増大します。単純な2方向分岐でも性能は半減しますが、より複雑な多方向分岐や入れ子になった条件文では、さらに劇的な性能低下が発生します。特に、ワープ内のスレッドが多数の異なる実行パスに分散する場合、実質的に単一スレッド実行に近い状態になってしまいます。
このため、効率的なGPUプログラミングでは、条件分岐を可能な限り回避し、すべてのスレッドが同一の実行パスを取るようアルゴリズムを設計することが不可欠です。条件分岐が避けられない場合は、ワープ境界を考慮して分岐を配置し、同一ワープ内での分岐を最小限に抑える戦略が重要になります。
14. 最適化技術2:低精度演算
14.1 FP32→FP16→INT8の精度削減によるオーダー単位の性能向上
最初にお話ししたいのは低精度についてです。これは大きなトリックです。これは重要なトリックです。常に行うべきです。Billyのこの図表に戻ると、ここには少し手品があります。これは本当に良く見えます。なぜなら数字がどんどん上がっているからです。
しかし、これらすべての年月にわたってGPUの進歩を実際に駆動しているものを見ると、実際には数値表現であることがわかります。FP32からFP16、INT8などに移行することで、GPU演算で桁違いの利得を得ることができます。
そして、なぜそれがそれほど重要なのかを明確にしましょう。計算しているものやウェイトなどのすべてにより少ないビットがある場合、移動するビットがはるかに少なくなります。たとえグローバルメモリからこれらのビットにアクセスしていても、それらははるかにはるかに懸念事項ではなくなります。
この精度削減による性能向上は、GPU コンピューティングにおける最も効果的な最適化手法の一つです。FP32(32ビット浮動小数点)からFP16(16ビット)、さらにINT8(8ビット整数)への移行により、メモリ使用量とメモリ帯域幅要求を劇的に削減できます。これは単なる数値的な改善ではなく、システム全体のアーキテクチャに根本的な影響を与える変化です。
Billyのチャートが示す驚異的な性能向上の「手品」は、実際にはこの数値表現の進歩によるものが大きな割合を占めています。同じハードウェア上でも、データ表現を変えるだけで桁違いの性能向上を実現できるのです。これは、計算ユニットの物理的改善だけでなく、データ効率の向上によって達成された成果なのです。
この手法の重要性は、現代の深層学習において証明されています。多くの実用的なモデルで、適切な精度管理により大幅な性能向上と電力効率の改善を同時に実現できることが実証されているのです。
14.2 算術強度の改善例:ReLU演算でのメモリアクセス半減効果
簡単な例を挙げて、単純な要素別演算の算術強度について考えてみましょう。x equals max zero and x を行い、サイズnのベクターに対してそれを行うとします。
単純に、これをfloat 32で行うとします。どれだけのメモリアクセスがありますか?xを読み取る必要があります。x less than zero の結果を書き込む必要があります。それはすべてfloat 32です。つまり、8バイトのようなものです。どれだけの演算を行いますか?x less than zero を行う必要があります。それは1つの比較演算です。そして1つのflopを行います。
つまり、単一の浮動小数点演算当たり8バイトを行います。これをfloat 16で行うと、floops強度はここで変わっていませんが、メモリアクセスを半分にしました。そして今、flop当たり4バイトになります。
ある意味で、float 16で実行できると仮定して、無料でメモリ帯域幅を2倍にしたようなものです。これが多くのものが設計される重要な部分です。
この具体例は、低精度演算の実用的な効果を明確に示しています。ReLU関数という単純な演算において、FP32では読み取り4バイト + 書き込み4バイト = 8バイトのメモリアクセスが必要ですが、FP16では読み取り2バイト + 書き込み2バイト = 4バイトに半減します。演算数は変わらないため、算術強度(演算数÷メモリアクセス量)が2倍に向上します。
この改善は「無料でメモリ帯域幅を2倍にした」と表現されているように、ハードウェアの物理的な改善なしに性能向上を実現できることを意味します。メモリがボトルネックとなっている多くのGPU ワークロードにおいて、この効果は極めて重要です。200-300クロックサイクルのグローバルメモリアクセス時間を考慮すると、転送データ量を半減できることの価値は計り知れません。
この原理は、Assignment の一部として混合精度や低精度訓練を試すことになる理由でもあります。実際のワークロードにおいて、どの程度の性能向上が達成できるかを体験することで、この最適化手法の実用的価値を理解できるのです。
14.3 混合精度行列乗算:16ビット入力、32ビット累積、32ビット出力
ここで重要な部分は、ネットワークと訓練アルゴリズムのすべての部分を低精度にするべきではないということです。行列乗算の例を挙げましょう。
混合精度の行列乗算では、入力を16ビットにします。これらは低精度です。そして、乗算を完全な32ビットで行います。これは、部分和を累積している中間計算を高精度にしたいからです。テンソルコアはFP32アキュムレータでこれを累積し、FP32結果を返します。もし望むなら、それを16ビットにダウンキャストすることができます。
つまり、入力は16ビットですが、累積のようなものは32ビットで行いたいかもしれません。16ビットストレージを使用できる演算があり、FP32やFP16に保持したい演算があります。
この混合精度アプローチは、数値安定性と性能の最適なバランスを実現するための戦略的な設計です。入力データを16ビットで保存することにより、メモリ使用量とメモリ帯域幅を大幅に削減できます。しかし、行列乗算の中間計算、特に部分和の累積では、数値精度の劣化を防ぐために32ビット精度を維持します。
テンソルコアのハードウェア設計は、まさにこの混合精度パターンに最適化されています。16ビット入力を受け取り、内部では32ビット精度で累積を行い、最終的に32ビット結果を出力します。この設計により、精度を犠牲にすることなく、メモリ効率と計算効率の両方を実現しています。
最終的な結果を再び16ビットにダウンキャストするかどうかは、アプリケーションの要件によって決まります。次の層の入力として使用する場合は16ビットにダウンキャストしてメモリ効率を維持し、高精度が必要な最終出力では32ビットのまま保持するという柔軟性があります。この段階的な精度管理こそが、現代の効率的な深層学習システムの基盤となっているのです。
14.4 演算別精度要件:範囲が必要な指数関数でのBF-16活用
多くの異なることがあります。16ビットストレージを使用できる演算があり、FP32やFP16に保持したい演算があります。動的範囲を持たない場合に爆発したりゼロになったりする可能性がある指数関数のような演算があるかもしれません。そして、それらをBF-16に入れたいかもしれません。
低精度で訓練されているときに、これらのモデルが実際に安定していることを確認するために行わなければならない多くの慎重なエンジニアリングがあります。しかし、それができるなら、それは本当に素晴らしいことです。なぜなら、メモリがボトルネックである場合、32ビットから16ビットに移行することで基本的にボトルネックのスループットを2倍にしたからです。
演算別の精度要件は、数値計算の特性を深く理解した上での戦略的判断を必要とします。指数関数のような演算では、入力値の微小な変化が出力に劇的な影響を与える可能性があるため、十分な動的範囲が不可欠です。FP16では表現可能な数値の範囲が限定されているため、指数関数の計算で容易にオーバーフローやアンダーフローが発生してしまいます。
BF-16(Brain Floating Point 16)は、この問題に対する解決策として設計されました。FP32と同じ指数部のビット数を維持することで、FP16よりもはるかに広い動的範囲を実現しています。仮数部の精度は FP16 より低くなりますが、指数関数のような範囲が重要な演算では、精度よりも範囲の方が決定的に重要なのです。
このような慎重なエンジニアリングは、現代の深層学習システムにおいて極めて重要です。各演算の数値特性を理解し、適切な精度を選択することで、数値安定性を保ちながら大幅な性能向上を実現できます。メモリがボトルネックとなっている現在のGPU 環境では、32ビットから16ビットへの移行により得られる2倍のスループット向上は、システム全体の性能に決定的な影響を与えるのです。
15. 最適化技術3:演算子融合(Operator Fusion)
15.1 工場の比喩:コンベアベルト(メモリ)がボトルネック
もう一つのことで、多くの人がCUDAカーネルを書くとか言うときに考えることだと思います。演算子融合は直感的でもあり、考えるのが楽しく自然なものでもあります。
GPUがどのように動作し、メモリがどのように動作するかについての一つの精神的モデルは、Horus Heathからのこの工場の楽しい図です。工場があり、工場が計算部分だと想像してください。小さなボックスウィジェットを取り込んで、小さな三角形ウィジェットを出力します。
計算を増やしても、メモリから計算へのベルトコンベア、つまりメモリから計算に取るものが有限の帯域幅である場合、第2の工場を使用することはできません。まだ、メモリから計算へのものを転送する速度によって制限されています。そこでボトルネックがあります。
もちろん、あなたはすでにそれを知っていました。私はメモリボトルネックのことをずっと打ち込んできました。しかし、実際に気づかずに大量のオーバーヘッドを被る可能性がある一つの陰湿な方法は、この左側の計算パターンの一種だと思います。
この工場の比喩は、GPU のメモリと計算の関係を非常に分かりやすく説明しています。工場(計算ユニット)がいくら高性能でも、原材料を供給するコンベアベルト(メモリ帯域幅)の速度が制限されていれば、全体の生産性はコンベアベルトの速度によって決まってしまいます。第2の工場を追加しても、コンベアベルトが1本しかなければ、その恩恵を受けることはできません。
この比喩の重要な洞察は、計算能力の向上だけでは性能向上に限界があるということです。現代のGPUまさにこの状況にあり、テンソルコアによる計算能力は飛躍的に向上しましたが、メモリ帯域幅の改善は相対的に緩やかです。結果として、多くのワークロードでメモリがボトルネックとなっています。
しかし、この比喩で最も重要なのは、「陰湿な方法でオーバーヘッドを被る」という警告です。プログラマーが意識しないうちに、非効率的なメモリアクセスパターンにより性能を大幅に損なってしまう可能性があるのです。特に、次に説明する左側の計算パターンは、一見合理的に見えても実際には深刻な性能問題を引き起こす典型例なのです。
15.2 ナイーブ実装vs融合実装:メモリアクセス回数の劇的削減
この左側の計算パターンの一種だと思います。この図の左側がメモリがある場所だと想像してください。右側が計算ユニットです。計算を行うために、正方形から始めて、正方形をメモリから計算に移します。何らかの操作を行います。それらを三角形に変えます。今、三角形をメモリに戻します。
そして、三角形が再び必要だと気づきます。そこで、それらを計算ユニットに戻します。今度は三角形が円になり、それからというように、計算をメモリに行ったり来たりと送り返します。これをかなりナイーブなアプローチと呼ぶかもしれません。
もしGPU上で演算をナイーブに行い、結果をグローバルメモリに直接戻すだけだった場合、これが最終的に得られるものです。データが行ったり来たりした回数を数えると、これはかなりひどいです。大量のメモリオーバーヘッドを被りました。
今、右側を見てください。この計算には依存関係がないので、正方形から三角形、円、長方形に行き、長方形を戻すことができるはずです。計算ユニット内ですべてを保持し続けることができます。それが右側の図で、これが融合カーネルの精神的モデルです。
データの系列で起こる一連の操作があり、ストレージに書き戻す代わりに、一箇所でできるだけ多くの計算を行い、メモリに戻す必要があるときだけ戻すということです。
この対比は、効率的なGPU プログラミングにおける最も重要な概念の一つを示しています。ナイーブな実装では、各演算の結果を律儀にグローバルメモリに書き戻し、次の演算で再度読み込むという無駄なサイクルを繰り返します。正方形→三角形→メモリ→三角形→円→メモリ→円→長方形というパターンでは、中間結果が何度もメモリを往復します。
200-300クロックサイクルのグローバルメモリアクセス時間を考慮すると、この無駄な往復は性能に壊滅的な影響を与えます。データが行ったり来たりする回数をカウントすると、実際の計算時間よりもメモリアクセス時間の方がはるかに長くなってしまいます。
融合実装では、依存関係のない一連の演算を計算ユニット内で連続して実行し、最終結果のみをメモリに書き戻します。正方形→三角形→円→長方形→メモリという単純なフローにより、メモリアクセスを劇的に削減できます。これは単なる最適化ではなく、GPU の物理的制約を理解した上での必然的な設計選択なのです。
15.3 実例:sin²x + cos²x計算での5カーネル→1カーネル融合
どのようにナイーブなコードを書けば、ナイーブな一連の起動を得ることができるかについて、非常に簡単な例があります。
ここに例があります。小さなニューラルネットワークモジュールを書いたとしましょう。sxを取り、sin²xとcos²xを生成するとします。簡単なコードです。これを実行すると、PyTorchの計算グラフはこのようになり、多くのCUDAカーネルを起動することになります。xを取り、sin xを計算するCUDAカーネルを起動し、cos xを計算するものを起動し、それからsin²xとcos²x、そしてsin²x + cos²xを起動します。
前に示した左側の図とまさに同じように、この計算を行うために多くの行ったり来たりが起こる必要があります。
しかし、もう少し賢くて、独自のCUDAカーネルを書いたか、torch compileのようなものを使った場合、これら5つの操作は実際にはあまり多くのメモリを使用しないことを簡単に認識できます。そして、グローバルメモリに物事を送り返すことなく、GPU上の単一スレッドですべてを行う単一の操作に融合することができます。
この実例は、演算子融合の実用的な効果を具体的に示しています。sin²x + cos²xという比較的単純な数学関数でも、ナイーブな実装では5つの個別のCUDAカーネルが起動されます:①sin(x)、②cos(x)、③sin²(x)、④cos²(x)、⑤加算。各カーネル間では中間結果がグローバルメモリに書き込まれ、次のカーネルで読み込まれるという無駄な往復が発生します。
各メモリアクセスに200-300クロックサイクルが必要であることを考えると、5つのカーネルによる計算よりも、4回のメモリ往復(8回のメモリアクセス)の方がはるかに多くの時間を消費してしまいます。実際の数学計算は数クロックサイクルで完了するのに、メモリアクセスで数千クロックサイクルを浪費しているのです。
融合実装では、すべての計算を単一スレッド内で連続して実行できます。入力xを読み込み、sin(x)とcos(x)を計算し、それぞれを2乗し、加算して結果を出力するまで、すべて高速な共有メモリまたはレジスタ内で完結します。これにより、グローバルメモリアクセスは入力読み込みと結果出力の2回のみに削減され、劇的な性能向上を実現できるのです。
15.4 Torch Compileによる自動融合の推奨
このような本当に簡単な融合操作は、コンパイラによって自動的に行うことができます。私はtorch compileを言及しました。もし、まだこれを行っていない場合は、どこでもtorch compileを使用することを強く検討すべきです。
Assignment 2でもtorch compileをお見せします。それはかなり素晴らしいものです。
Torch Compileの推奨は、現代のPyTorch開発において極めて実用的なアドバイスです。手動でCUDAカーネルを書いて演算子融合を実装するのは高度な技術と深い理解を必要としますが、torch compileを使用することで、これらの最適化を自動的に適用できます。
torch compileは、PyTorchの計算グラフを分析し、融合可能な演算を特定して、自動的に最適化されたカーネルを生成します。先ほどのsin²x + cos²x の例のような単純なケースから、より複雑な演算チェーンまで、幅広い融合パターンを認識できます。これにより、開発者は低レベルの最適化に時間を費やすことなく、効率的なGPUコードを実行できます。
「どこでもtorch compileを使用する」という強い推奨は、その効果の確実性を示しています。多くの場合、コードを1行追加するだけで大幅な性能向上を得ることができます。Assignment 2でその実演を見ることで、この最適化技術の実用的な価値を体験できるでしょう。
torch compileは、GPU最適化の民主化を実現するツールでもあります。従来は専門知識を必要とした高度な最適化を、一般の開発者でも簡単に利用できるようにしており、効率的なGPUプログラミングの普及に大きく貢献しているのです。
16. 最適化技術4:再計算(Recomputation)
16.1 計算リソースとメモリ帯域幅のトレードオフ
私たちができるもう一つのことは再計算と呼ばれます。再計算は、メモリアクセスを行う必要を避けるために、より多くの計算を費やすというアイデアです。
元の逆伝播の講義を思い出してください。実際、これはCS221からのものです。何をしますか?一番下にある入力を取ります。これらは黄色のものです。そして、活性化を上向きに伝播します。これらも木の上の黄色の値です。そして、ヤコビアンを逆向きに計算します。これらはエッジ上の緑の値です。そして勾配を計算するために、ヤコビアンと活性化を掛け合わせるようなことをします。勾配を逆向きに伝播します。
もし考えてみると、順伝播の後、これらの黄色の値は保存されなければなりません。そして、それらは保存され、それらを保存したグローバルメモリから取り出して計算ユニットに入れる必要があります。しかし、それは実際に大量のメモリの入出力が起こることかもしれません。代わりに、これを実際に避けることができるかもしれません。
この再計算のアプローチは、現代のGPUにおける根本的なリソーストレードオフを活用しています。以前に説明したように、計算能力(灰色の線)は1から100,000倍という驚異的な成長を遂げた一方、メモリ帯域幅(緑の線)の成長は相対的に緩やかでした。この結果、多くのワークロードでメモリがボトルネックとなり、計算ユニットが十分に活用されていない状況が生まれています。
再計算は、この不均衡を戦略的に利用する最適化手法です。余剰となった計算能力を活用して、メモリアクセスを削減するのです。従来の逆伝播では、順伝播で計算された活性化値(黄色の値)をメモリに保存し、逆伝播時に再度読み込む必要がありました。しかし、これらの値を保存せずに、必要な時に再計算することで、メモリアクセスを大幅に削減できます。
この戦略が効果的な理由は、200-300クロックサイクルのグローバルメモリアクセス時間と比較して、多くの計算は数クロックサイクルで完了するからです。メモリアクセス待機で計算ユニットがアイドル状態になるくらいなら、その時間を使って値を再計算する方がはるかに効率的なのです。
16.2 3重シグモイド関数の実例分析
再計算がどのように物事を高速化できるかの例を挙げましょう。ここに、私が書くかもしれないもう一つの愚かな関数があります。3つのシグモイドを互いの上に積み重ねるだけです。左を見ることができます。それは順伝播グラフです。それは、互いの上に積み重ねられた3つのシグモイドのあなたの精神的モデルであるはずです。
この計算グラフでは、シグモイドを計算し、S1とS2という活性化を保存し、出力があります。それが私の順伝播です。次に、この逆伝播は非常にひどいです。逆伝播を行うとき、S1とS2を取る必要があり、このoutボックスに逆向きに入ってくる勾配を取る必要があり、それをこの逆計算にプッシュして、xの勾配を得ます。
順伝播では、xの1回のメモリ読み取りと、S1、S2、outの3回のメモリ書き込みを行う必要があります。逆伝播では、S1、S2、outの3回のメモリ読み取りと、勾配計算後の1回のメモリ書き込みが必要です。つまり、8回のメモリアクセスを行う必要があり、行列乗算が全くないため、算術強度が非常に低いです。
この3重シグモイド関数の分析は、従来の逆伝播における重大な非効率性を明確に示しています。一見単純な関数でも、ナイーブな実装では大量のメモリアクセスが発生してしまいます。
順伝播段階では、入力xを読み込み(1回のメモリ読み取り)、3つのシグモイド演算を順次実行し、中間結果S1、S2と最終出力outをグローバルメモリに保存します(3回のメモリ書き込み)。各メモリアクセスに200-300クロックサイクルが必要であることを考えると、実際のシグモイド計算時間よりもメモリアクセス時間の方がはるかに長くなります。
逆伝播段階では状況がさらに悪化します。保存された活性化S1、S2、outを読み込み(3回のメモリ読み取り)、上位層からの勾配信号を使って逆伝播計算を実行し、最終的に入力xに対する勾配を出力します(1回のメモリ書き込み)。
合計8回のメモリアクセスが必要で、しかも行列乗算のような高い算術強度を持つ演算が含まれていないため、この関数は典型的なメモリバウンド問題となっています。計算ユニットは大部分の時間をメモリアクセス待機で費やし、実際の有用な計算はわずかな時間しか行われていないのです。
16.3 メモリアクセス削減:8回→5回(37.5%削減)
再計算のアイデアは、これらの活性化を全く保存したくないということです。逆伝播で、オンザフライで再計算するだけです。
新しい順伝播では、S1とS2を保存しません。xを入力として取り、シグモイドを計算し、出力を得ます。そこで、xに対する1回のメモリ読み取り、outに対する1回のメモリ書き込みがあります。
次に、私の逆伝播では、もう活性化がありません。そこで行うことは、上から入ってくる逆向きの信号であるd outと、私の入力であるxの両方を取ることです。それらの2つを取り、これは2回のメモリ読み取りです。そして、私のSM、私のローカルメモリで、オンザフライで、これらのシグモイドのそれぞれを計算し、それらを逆伝播グラフに入れます。S1、S2、outをローカルメモリ内でオンザフライで再計算します。
ここではグローバルメモリ読み取りが起こっていないため、そして1回のメモリ書き込みがあり、それはdxです。そこで、2つを比較すると、まったく同じ計算に対して8分の5のメモリアクセスがあります。3つのシグモイドを再計算する必要がありますが、とにかくメモリ制限で実行していた場合、これは素晴らしいトレードオフです。計算が少なすぎたものを、メモリ帯域幅が少なすぎたものと交換したからです。
この37.5%のメモリアクセス削減は、再計算手法の具体的な効果を数値で示しています。8回から5回への削減は単なる改善ではなく、システム全体の性能特性を根本的に変える可能性があります。
新しい順伝播では、中間活性化の保存を完全に省略します。入力x の読み込み(1回)と最終出力outの書き込み(1回)のみで、合計2回のメモリアクセスに削減されます。これは従来の4回(1読み込み + 3書き込み)から50%の削減です。
逆伝播では、保存された活性化の代わりに、入力xと上位層からの勾配d_outを読み込みます(2回の読み取り)。重要なのは、S1、S2、outの再計算がSMのローカルメモリ内で実行されることです。これらの中間値はレジスタや共有メモリに保持され、20クロックサイクル程度の高速アクセスで処理されます。最終的にdxを出力(1回の書き込み)して、逆伝播段階は合計3回のメモリアクセスで完了します。
この戦略が「素晴らしいトレードオフ」と評価される理由は、余剰計算能力を不足メモリ帯域幅と交換しているからです。3つのシグモイド関数の再計算に必要な時間よりも、3回のグローバルメモリアクセスを回避することで節約される時間の方がはるかに大きいのです。メモリ制限環境では、この交換により全体の実行時間が大幅に短縮されます。
16.4 勾配チェックポイントとの違い:速度向上 vs メモリ節約
もちろん、これは異なります。これは、メモリ節約のための活性化の勾配チェックポイントと再計算と同じトリックです。しかし、これは異なる理由で行われています。これは実行速度のためであり、メモリ不足だからではありません。それは同じ技術ですが、異なる目標のためです。
この区別は、最適化技術の応用における重要な洞察を提供しています。勾配チェックポイントと再計算は、表面的には同じ技術的手法を使用しています。どちらも中間活性化の保存を省略し、必要時に再計算を行います。しかし、その動機と目標は根本的に異なります。
従来の勾配チェックポイントは、主にメモリ制約の解決を目的としています。大規模なニューラルネットワークでは、すべての活性化を保存するために必要なメモリ量が利用可能なメモリを超えてしまうことがあります。この場合、メモリ不足によるトレーニングの停止を回避するために、一部の活性化のみを保存し、他は再計算することでメモリ使用量を削減します。計算時間の増加は、トレーニングを継続するための必要なコストとして受け入れられます。
一方、ここで説明している再計算は、実行速度の向上を主目的としています。メモリが不足しているわけではなく、メモリアクセスがボトルネックとなっている状況で、意図的に再計算を選択します。余剰となった計算能力を活用して、遅いグローバルメモリアクセスを回避することで、全体の実行時間を短縮します。
この同じ技術の異なる応用は、GPU最適化における文脈の重要性を示しています。ハードウェアの特性、ワークロードの性質、制約条件によって、同じ技術でも全く異なる価値を提供するのです。現代のGPU環境では、メモリ帯域幅がボトルネックとなることが多いため、速度向上を目的とした再計算の重要性がますます高まっているのです。
17. 最適化技術5:メモリ合体(Memory Coalescing)
17.1 DRAMのバーストモードメカニズム
これは、私がGPUとDRAMがどのように動作するかのハードウェアモデルを本当に調べ始めるまで知らなかった、本当に興味深いものの一つだと思います。
遅いメモリ、グローバルメモリと呼ばれるGPU内のDRAMは、実際には非常に非常に遅いです。それをより速くするために、ハードウェアレベルで行われている特定の最適化があります。
DRAMに対してハードウェアレベルで行われる最適化の一つは、メモリの一部を読み取りに行くとき、実際にはその値だけを返すのではないということです。実際には、メモリの全体のチャンクを返します。これはバーストモードと呼ばれます。
この大きなメモリブロックの最初の値を読み取ろうとしたとしましょう。メモリが0だけを返す代わりに、実際には0、1、2、3を返します。4つの値を一度に返します。「ここにあります。将来、1、2、3も必要になると確信しています」のような感じです。
各アドレス空間はバーストセクションと呼ばれるものに分割されます。そして、探していたものだけではなく、バーストセクション全体が与えられます。
このDRAMのバーストモードメカニズムは、メモリアクセスの物理的制約に基づく巧妙な最適化です。単一のメモリセルにアクセスするためには、まずアドレスを指定し、そのアドレスに対応するメモリセルからデータを読み出す必要があります。この過程で最も時間がかかるのは、実際にはメモリセルの場所を特定し、アクセス準備を整える初期段階です。
一度この初期段階が完了すると、隣接するメモリセルから追加のデータを読み出すコストは相対的に小さくなります。バーストモードは、この特性を活用して、単一のアクセス要求に対して連続する複数の値を返します。例えば、アドレス0のデータを要求すると、0、1、2、3の4つの値がまとめて返されます。
この設計の背景には、空間局所性の原理があります。プログラムが特定のメモリ位置にアクセスした場合、近接するメモリ位置にも短時間内にアクセスする可能性が高いという統計的傾向に基づいています。DRAMハードウェアは、この傾向を予測して、要求されていない隣接データも事前に提供することで、将来のアクセス要求に対する応答時間を短縮しているのです。
バーストセクションという概念により、メモリ空間は固定サイズのブロックに分割され、各ブロック内のデータは効率的に一括取得できるようになっています。この仕組みを理解し活用することが、効率的なGPUプログラミングの重要な要素なのです。
17.2 ハードウェア理由:アンプへの信号移動が高コスト、その後の複数バイト取得は低コスト
なぜメモリが要求した1つだけではなく、3つの余分なバイトを無料で提供するのか、非常に神秘的に思えるかもしれません。
非常に興味深いハードウェア理由があります。メモリにアドレシングする際、それらのバイトをアンプに移動させるために、メモリから信号を送出する必要があります。それが遅いステップです。それを行った後は、多くの多くのバイトを無料で取得できます。
そのため、このバーストセクションの存在理由は、実際にデータが保存されている場所からこのアンプにデータを移動するより高価なステップをマスクしているのです。しかし、それとは関係なく、もし良いメモリアクセスパターンがあれば、メモリアクセスを大幅に加速できるかもしれません。
このハードウェアレベルの詳細説明は、バーストモードの物理的基盤を明確にしています。DRAM の内部構造では、実際のデータはメモリセルアレイに保存されていますが、これらの微小な電荷を読み取るためには、まず信号をセンスアンプ(増幅器)に移動させる必要があります。
このセンスアンプへの信号移動が、メモリアクセスにおける律速段階となっています。メモリセルからアンプまでの物理的距離と、微小な電荷を増幅可能なレベルまで引き上げる処理に時間がかかるのです。しかし、一度センスアンプにデータが到達すると、そこから連続する複数のバイトを読み出すことは相対的に高速で実行できます。
この物理的制約こそが、「3つの余分なバイトを無料で提供する」理由です。最初の1バイトを読み出すために必要な高コストなセットアップ処理を、隣接する複数バイトで償却できるからです。アンプへの信号移動という「高価なステップをマスク」することで、全体的なメモリスループットを向上させているのです。
この理解は実用的な意味を持ちます。プログラマーがメモリアクセスパターンを設計する際、この物理的特性を活用できれば、同じハードウェアでもメモリアクセス性能を劇的に改善できます。バーストセクション境界を意識したデータ配置とアクセスパターンにより、ハードウェアの投資なしに大幅な性能向上を実現できるのです。
17.3 バーストセクション最適化:4倍スループット向上の可能性
もし良いメモリアクセスパターンがあれば、メモリアクセスを大幅に加速できるかもしれません。
この全体のブロックを読み取りたい場合、それをランダムな順序でアクセスすると、基本的にクエリの長さに大体等しい回数クエリする必要があります。
しかし、行って最初の値をチェックした場合、この全体のバーストセクションを一度に取得します。そして、4番目の値をチェックしに行くと、2番目のバーストセクション、2番目のバーストセクションを一度に取得します。
そのため、メモリアクセスが本当に賢く、各バーストセクションから必要なビットだけをアクセスする場合、基本的に4倍のスループットを得ることができます。
この4倍スループット向上の可能性は、メモリアクセスパターン最適化の重要性を具体的に示しています。同じデータを読み取る場合でも、アクセス順序によって性能が劇的に変化するのです。
ランダムアクセスの場合、各メモリ要求が異なるバーストセクションを対象とすると、それぞれに対して完全なメモリアクセス手順(アドレシング、センスアンプへの信号移動、増幅)を実行する必要があります。4つの値を読み取るために4回の完全なメモリアクセスが必要となり、各アクセスに200-300クロックサイクルが消費されます。
対照的に、シーケンシャルアクセスでは、最初の値(例えば0番目)にアクセスすることで、バーストセクション全体(0、1、2、3)が一度に取得されます。次に必要な値(例えば4番目)にアクセスすると、次のバーストセクション(4、5、6、7)が取得されます。この方法では、4つの値を読み取るために必要な完全なメモリアクセスは2回だけです。
「各バーストセクションから必要なビットだけをアクセス」することで、実質的にメモリアクセス回数を4分の1に削減できます。これは、同じハードウェア、同じデータ量でありながら、アクセスパターンの最適化だけで4倍のスループットを実現できることを意味します。この原理こそが、効率的なGPUプログラミングにおいてメモリレイアウトとアクセスパターンの設計が決定的に重要である理由なのです。
17.4 行列乗算での実例:行走査vs列走査のパフォーマンス差
これはメモリ合体と呼ばれます。ワープ内のすべてのスレッドが同じバースト内に収まる場合、基本的にスマートハードウェアとプログラミングモデルは、これらのクエリをグループ化します。0、1、2、3をクエリする代わりに、それらをグループ化して「0を与えて」と言い、この種のバーストモードDRAMから0、1、2、3をすべて一度に読み出すことができます。
ワープは32の番号付きスレッドであることを覚えておいてください。これらのワープからのメモリアクセスは一緒に起こります。これらのワープがこれらの種類のバーストセクションに読み込んでいるとき、1つずつ個別に1、2、3、4バイトを取得するのではなく、一度にすべての4バイトを取得できるように最適化が行うことができます。これにより、メモリのスループットが4倍になります。
これらは非常に単純なことですが、実際には非常に重要です。行列乗算を行うと想像してください。これは、本当にゼロからCUDAで何かのニューラルネットワークを実装しようとした場合に行う必要がある中核的なことです。
この場合、私の行列を2つの方法のうちの1つで読み取ろうと想像してください。行を横断して読み取ることができます。各スレッドが行を横断します。または、列順で読み取ることができます。各スレッドが列を下に行きます。
この左側のもの、異なる行を横断しているところ、各スレッドが異なる列にアクセスしているこの左側のモデルはかなり遅くなります。メモリ読み取りが合体されないからです。
一方、これらのスレッドが下に行く右側、つまり行が増加するところでは、これらのメモリ読み取りが合体されます。
この行列乗算での実例は、メモリ合体の概念を具体的なプログラミング文脈で説明しています。32個のスレッドで構成されるワープが同時にメモリアクセスを行う際、そのアクセスパターンが性能に決定的な影響を与えるのです。
左側の行走査パターンでは、各スレッドが異なる行の要素にアクセスします。例えば、時刻t=1において、スレッド0が行0の要素、スレッド1が行1の要素、スレッド2が行2の要素にアクセスします。これらの要素は行列内で物理的に離れた位置にあるため、異なるバーストセクションに属します。結果として、32個のスレッドが32個の異なるバーストセクションにアクセスすることになり、32回の完全なメモリアクセスが必要になります。
対照的に、右側の列走査パターンでは、すべてのスレッドが同一行内の連続する要素にアクセスします。時刻t=1において、全32スレッドが同じ行の連続する32要素にアクセスするため、これらの要素は同じまたは隣接するバーストセクションに属します。ハードウェアはこれらのアクセスを自動的に合体し、数回のバーストアクセスで全てのデータを取得できます。
この違いにより、右側のパターンは左側と比較して数倍から数十倍の性能向上を実現できます。同じ計算を行うにも関わらず、メモリアクセスパターンの違いだけで、これほど劇的な性能差が生じるのです。この原理を理解することは、効率的な行列演算ライブラリやCUDAカーネルを実装する上で不可欠な知識なのです。
18. 最適化技術6:タイリング(Tiling)
18.1 グローバルメモリアクセス最小化の基本概念
そして、最後の大きなものに到達します。これはタイリングのアイデアです。タイリングは、実行する必要があるグローバルメモリアクセスの量を最小化するために、メモリアクセスをグループ化したいというアイデアです。
この1つを説明するために、行列乗算の例を通して説明しようと思います。うまくいけば、なぜナイーブなアルゴリズムで行列乗算を行うことが非常に問題があるのかを説明できるでしょう。その後、同じアイデアのタイル化されたバージョンをお見せします。うまくいけば、実行する必要があるグローバルメモリ読み取りの数を減らす理由を見ることができるでしょう。
理想的には、一定時間をかけてグローバルメモリから共有メモリに断片を読み込み、共有メモリで大量の計算を行い、そのデータの断片で完了したいと思います。それが理想的な結果です。グローバルメモリアクセスを最小化しました。
タイリングの基本概念は、GPU の階層的メモリ構造を戦略的に活用する最適化手法です。グローバルメモリ(200-300クロックサイクル)と共有メモリ(20クロックサイクル)の間の圧倒的な速度差を考慮すると、可能な限りグローバルメモリアクセスを削減し、高速な共有メモリでの計算を最大化することが性能向上の鍵となります。
この手法の核心は「グループ化」の概念にあります。データを小さなタイル(瓦)状のブロックに分割し、各タイルを一度に共有メモリに読み込んで、そのタイル内で可能な限り多くの計算を実行します。これにより、同じデータを複数回グローバルメモリから読み込む必要がなくなり、メモリアクセス効率が劇的に改善されます。
理想的な実行パターンは明確です:①グローバルメモリから共有メモリへの一括データ転送、②共有メメリでの集約的計算処理、③次のタイルへの移行。このサイクルにより、高速な計算ユニットが常に有用な作業を行い続け、遅いグローバルメモリアクセス待機によるアイドル時間を最小限に抑えることができます。
この概念は、現代のGPU プログラミングにおいて最も重要な最適化技術の一つであり、後に説明するFlash Attentionのような先進的なアルゴリズムの基盤ともなっているのです。
18.2 ナイーブ行列乗算の問題:非合体アクセスと重複読み込み
この非常に単純な行列乗算アルゴリズムから始めましょう。左側にMマトリックスがあります。上にNマトリックスがあります。行列行列積を計算するために、Mの行とNの列を横断し、内積を取り、対応する行のPマトリックスに保存する必要があります。
ここに各スレッドを書き出しました。スレッド0 1、1 0、1...出力を保存している場所に対応し、個々の要素にアクセスする順序のアクセス順序です。
ここで注目してください。メモリアクセスが合体されていないことです。行マトリックスは非合体順序でアクセスされ、重複したメモリアクセスがあります。m[0,0]が最初のスレッドでアクセスされ、m[0,0]がここでアクセスされ、n[0,0] n[1,0]が2つの異なるスレッドでアクセスされています。
これらの値は、グローバルメモリから多くの異なるスレッドに何度も何度も読み込まれています。これは潜在的に非常に遅くなる可能性があります。
このナイーブ実装の分析は、効率的でない行列乗算における2つの深刻な問題を明確に示しています。まず、非合体メモリアクセスの問題です。各スレッドが異なる行や列の要素にアクセスするため、これらの要素は異なるバーストセクションに配置されている可能性が高く、メモリ合体の恩恵を受けることができません。
例えば、行列Mの行方向アクセスでは、各スレッドが異なる行の要素を読み取るため、物理的に離れたメモリ位置への個別アクセスが発生します。前述のバーストモードの説明で見たように、こうしたランダムアクセスパターンでは、各要素に対して完全なメモリアクセス手順を実行する必要があり、4倍のスループット向上の機会を逸してしまいます。
さらに深刻なのは重複読み込みの問題です。m[0,0]が複数のスレッドで読み込まれ、n[0,0]やn[1,0]も複数のスレッドでアクセスされています。これは、同じデータが200-300クロックサイクルの高コストでグローバルメモリから繰り返し読み込まれることを意味します。
この非効率性は、32個のスレッドからなるワープが同時に実行される環境では特に問題となります。多数のスレッドが同じデータに対して独立したグローバルメモリアクセスを行うため、メモリ帯域幅が無駄に消費され、計算ユニットは大部分の時間をメモリ待機で費やすことになるのです。
18.3 タイル化行列乗算アルゴリズム:2×2タイル例での詳細解説
では、グローバルメモリの読み書きが多すぎることを避けることができるかという問題があります。私が理想的にやりたいことを最初に説明させてください。理想的な結果は、一定時間をかけてグローバルメモリから共有メモリに断片を読み込み、共有メモリで大量の計算を行い、そのデータの断片で完了したいということです。
では、この行列乗算の世界でこれをどのように行うことができるでしょうか?
今、私が行おうとしていることは、マトリックスMとマトリックスNの両方を取り、それらを切り分けることです。タイルに切り分けます。ここで2×2タイルに切り分けました。2×2のMタイルと2×2のNタイルがあります。
基本的に、各行列内の小さなサブマトリックスがあります。今、私の共有メモリが、各SM内でこれらのサブマトリックスを収めるのに十分大きいと想像してください。
これにより、計算を行うことができる非常に非常に単純なアルゴリズムが得られます。最初に、ここの左上のm00タイルと、共有メモリにN00タイルも読み込みます。
これで計算できる部分和があります。m00のm01とn0 n0の行積を取り、それをp0に増分できます。ここでさまざまなサブマトリックスで埋めることができるすべての異なるサブマトリックスでも同じことができます。
これら2つのタイルの処理を完全に終えたら、ここに新しいタイルを読み込むことができます。MタイルとN2.0タイルを共有メモリに読み込んで、その計算をMタイルと繰り返し、Pの部分和を増分することができます。
これで、実行する必要があるグローバルメモリアクセスの量を本当に統合し、削減しました。一度にできるだけ多くのメモリを共有メモリに読み込み、そのタイルで実行できるすべてのサブマトリックス計算を行い、次の計算に移ります。
このタイル化アルゴリズムの詳細解説は、効率的な行列乗算の核心的な戦略を示しています。2×2という小さなタイルサイズを使用することで、概念を明確に理解できます。
まず、元の大きな行列MとNを、管理可能な小さなサブマトリックス(タイル)に分割します。各2×2タイルは、共有メモリの制限内に収まるサイズに設計されています。これにより、タイル全体を一度に高速な共有メモリに読み込むことが可能になります。
アルゴリズムの実行は段階的に進行します。第1段階では、M行列の左上タイル(m00)とN行列の対応するタイル(N00)を共有メモリに読み込みます。この時点で、グローバルメモリアクセスは2回(2つのタイル読み込み)のみです。
共有メモリに読み込まれたタイル間で、すべての可能な部分積計算を実行します。m00タイルの行とN00タイルの列の内積を計算し、結果行列Pの対応する位置に累積します。これらの計算はすべて20クロックサイクルの高速共有メモリアクセスで実行されます。
第1タイルペアの処理が完了すると、次のタイルペア(MタイルとN20タイル)を読み込み、同様の計算を繰り返します。この段階的な処理により、各データ要素を一度だけグローバルメモリから読み込み、共有メモリ内で最大限活用することができます。
重要なのは、タイル読み込み時にメモリ合体も最適化できることです。タイル内の要素は連続配置されているため、列順序または行順序での効率的な読み込みが可能になり、前述のバーストモード最適化も同時に活用できるのです。
18.4 タイリング数学:グローバルメモリ読み込みのt分の1削減効果
少しタイリング数学を行うことができます。行列A、行列B、行列Cがあるとしましょう。完全な行列、これらは正方行列がサイズNであるとしましょう。サイズTのタイルがあるとしましょう。
非タイル化行列乗算を行う場合、行と列を調べているだけなら、すべての入力はグローバルメモリからn回読み取られます。各入力はグローバルメモリからn回読み取られるため、これらのそれぞれはn回読み取られます。
タイル化行列乗算を行う場合、グローバル読み取りはタイル上で動作しています。各入力をグローバルメモリからn/t回読み取り、各タイル内でt回読み取ります。もちろん、行わなければならない読み取りの総数を減らすことはできません。すべての行列要素を読み取る必要がありますが、高速共有メモリに読み取りをシフトできます。グローバルメモリからn/t回読み取り、共有メモリからt回読み取ります。
もし大きなタイルを格納できる大きな共有メモリがある場合、それはグローバルメモリから来る必要があるデータの総量でt倍の削減です。
このタイリング数学は、最適化効果を定量的に示す重要な分析です。サイズN×Nの正方行列乗算において、各要素が複数回アクセスされる回数の変化を数式で表現できます。
非タイル化実装では、行列Cの各要素C[i,j]を計算するために、行列Aの行i全体(N個の要素)と行列Bの列j全体(N個の要素)が必要です。N×Nの結果行列を完成させるためには、各入力要素が平均してN回グローバルメモリから読み取られることになります。これは、同じデータが繰り返し200-300クロックサイクルの高コストでアクセスされることを意味します。
タイル化実装では、この読み取りパターンが根本的に変化します。各入力要素のグローバルメモリからの読み取り回数がn/t回に削減されます。例えば、N=1024、T=32の場合、グローバル読み取り回数は1024回から32回(1024/32)へと約32分の1に削減されます。
削減された読み取り回数の差(n - n/t = n(t-1)/t回)は、高速な共有メモリアクセスに置き換えられます。共有メモリでのt回読み取りは、20クロックサイクル×T回 = 20T クロックサイクルで完了しますが、削減されたグローバルメモリアクセスは約200-300×n×(t-1)/t クロックサイクルの節約を意味します。
「t倍の削減」という表現は、大きなタイルサイズTを使用できる場合の理論的最大効果を示しています。T=64の場合、理論的には64倍のグローバルメモリアクセス削減が可能で、これは数十倍から数百倍の性能向上に直結する可能性があるのです。
19. タイリングの複雑性
19.1 離散化問題:257サイズ行列での6タイル必要性とSMアイドル化
タイリングは非常に複雑です。これは、GPU と行列乗算性能について多くの多くの混乱の源です。
発生する可能性がある1つのことは、離散化について質問し始めることです。タイルサイズが128であると想像してください。それは良い丸いタイルサイズのように思えます。
しかし、完全な行列が256サイズである場合、それは素晴らしいです。それは2×2タイルです。物事がうまく読み込まれます。今、列側で257サイズのタイルがあるとしましょう。今、これは悪い時間です。なぜなら、この行列をカバーするために6つのタイルが必要だからです。そして、右側の2つのタイルは非常に、非常にスパースです。そこにはあまり多くのものがありません。
この問題は、各タイルがSMに割り当てられることです。これらのタイルのそれぞれがブロックになり、各スレッドが各タイル内で動作することになります。そのため、右側の2つのタイル、それらはあまり多くのことを行わないでしょう。それらのSMは基本的にアイドル状態になります。
計算制限されていた場合、SM間でより均等に負荷を分散したかったでしょう。
この離散化問題は、タイリング最適化における最も微妙で実用的な課題の一つです。理論的には完璧に見えるタイリング戦略が、実際の行列サイズとの相互作用により深刻な性能問題を引き起こす可能性があるのです。
128×128という「良い丸いタイルサイズ」は、多くの理由で魅力的です。2の累乗であり、メモリアライメントに適し、多くのGPUアーキテクチャの特性に合致します。256×256行列では、正確に2×2=4個のタイルで完全に分割でき、各SMが均等な作業量を受け取ります。
しかし、行列サイズが257×256に変わった瞬間、状況は劇的に悪化します。列方向で257要素をカバーするには、128要素のタイルでは不十分で、3個のタイルが必要になります(128+128+1)。結果として、2×3=6個のタイルが必要となり、右側の2個のタイルは1列分(257-256=1要素)のデータしか含まず、99.6%が空の状態になります。
この問題の深刻さは、GPUの実行モデルに起因します。各タイルは独立したブロックとしてSMに割り当てられ、そのSM内の全スレッドがそのタイル内で作業します。右側のスパースなタイルを担当するSMは、ほとんど何もすることがなく、アイドル状態で待機することになります。
計算制限されたワークロードでは、この不均等な負荷分散により全体性能が大幅に低下します。4個のSMが効率的に動作している間、2個のSMがほぼ完全にアイドル状態となるため、実質的にGPUリソースの3分の1が無駄になってしまうのです。このような「境界効果」は、実用的なアプリケーションにおいて予想外の性能低下を引き起こす典型的な例なのです。
19.2 タイルサイズ決定要因:メモリ合体、共有メモリ制限、均等分割
そのため、この種のシナリオを避けるためにタイルサイズを最適化する必要があります。しかし、実際には、タイルサイズの設定には多くの複雑なことが関わっています。
メモリアクセスを合体させることについて慎重に考える必要があります。共有メモリサイズを超えてはいけません。タイルが大きすぎてはいけません。そして、SM間でできるだけ均等に負荷を分散できるように、行列次元をうまく均等に分割したいと思います。この最後にある種の十分に活用されていないSMのような状況にならないようにしたいのです。
タイルサイズの決定は、複数の制約条件を同時に満たす複雑な最適化問題です。各要因が相互に影響し合うため、単純な解決策は存在しません。
まず、メモリ合体の要件があります。前述のバーストモード最適化を活用するため、タイルサイズは連続メモリアクセスを促進する必要があります。32スレッドのワープが同時にアクセスする要素が、同じまたは隣接するバーストセクション内に配置されるよう、タイル境界を慎重に設計する必要があります。これは、タイルサイズが特定の2の累乗の倍数であることを示唆することが多いです。
次に、共有メモリ制限という物理的制約があります。各SMの共有メモリ容量は有限であり、通常48KB-164KB程度です。タイルサイズが大きすぎると、必要なデータを共有メモリに収容できなくなり、タイリングの恩恵を受けることができません。さらに、入力タイル2個分(AタイルとBタイル)プラス出力用のワーキングスペースを考慮すると、実際に使用可能なタイルサイズはさらに制限されます。
最も複雑なのは均等分割の要件です。タイルサイズは、実際の行列次元を可能な限り均等に分割できるよう選択する必要があります。理想的には、各次元がタイルサイズの整数倍となることが望ましいですが、実用的なアプリケーションでは任意のサイズの行列を扱う必要があります。このため、一般的な行列サイズに対して最適化されたタイルサイズを選択するか、動的にタイルサイズを調整する仕組みが必要になります。
これらの制約条件すべてを満たす最適なタイルサイズを見つけることは、GPUプログラミングにおける最も困難な課題の一つであり、多くの場合、経験的な調整と性能測定による反復的最適化が必要になるのです。
19.3 バーストセクションとの相互作用:アライメント不良による2倍メモリアクセス
もう一つの非常に非常に複雑なことは、タイリングとバーストセクションの相互作用です。
このような行列レイアウトがあると想像してください。私の素晴らしいバーストセクションがあります。各バーストセクションがタイルときれいに並んでいます。このタイルを読み取るために、4つの異なるバーストセクションを取得するだけで、この全体のタイルを取得しました。
今、最後に1つの余分な要素を追加し、行列のレイアウト方法で、バーストセクションが溢れるようになったと想像してください。今、起こっていることは、このタイルを読み込むとき、この最初の部分を読み込み、それは本当に素晴らしいです。最初の行全体をバーストセクションとして取得します。今、2番目の行では、これは実際に2つの異なるバーストセクションに属しています。
この2番目の行を取得するために2つの読み取りを行う必要があり、以下同様です。そのため、1つの余分な要素を最後に追加したことで、バーストセクションとタイルレイアウトのアライメントを崩し、基本的にメモリアクセス数を2倍にしました。
このバーストセクションとタイリングの相互作用は、GPU最適化における最も微妙で見落としやすい問題の一つです。表面的には些細に見える1要素の追加が、システム全体の性能を劇的に悪化させる可能性があるのです。
理想的なアライメント状態では、タイル境界がバーストセクション境界と完全に一致します。各タイル行が正確にバーストセクション境界で始まり終わるため、行全体を単一のバーストアクセスで効率的に読み取れます。4つのバーストアクセスで完全な2×2タイル(仮定)を取得でき、メモリ帯域幅を最大限活用できます。
しかし、たった1要素の追加により、この繊細なアライメントが崩壊します。行列の物理的レイアウトがシフトし、タイル内の各行が2つのバーストセクションにまたがるようになります。最初の行の一部は1つのバーストセクションの末尾に、残りの部分は次のバーストセクションの先頭に配置されます。
この結果、各行を読み取るために2回のバーストアクセスが必要になります。4行のタイルでは、必要なバーストアクセス数が4回から8回に倍増します。実際の有用データは同じでありながら、メモリアクセス時間が2倍になるのです。200-300クロックサイクルのグローバルメモリアクセス時間を考慮すると、この2倍の増加は全体性能に壊滅的な影響を与えます。
この問題の教訓は、効率的なGPUプログラミングでは、アルゴリズムレベルの最適化だけでなく、データレイアウトとハードウェア特性の相互作用まで考慮した包括的な設計が必要だということです。
19.4 パディングによる解決策
この回避方法は、タイルやバーストセクションと整列するように、きれいな丸い行列サイズを得るためにパディングを行う必要があることです。
パディングによる解決策は、アライメント問題に対する実用的なアプローチです。行列の実際のサイズが最適でない場合、人工的に要素を追加して理想的なサイズに調整するのです。
例えば、257×256の行列の場合、列次元を260や256に調整することでアライメントを回復できます。追加された要素にはゼロを埋めるか、計算に影響しない値を設定します。これにより、バーストセクション境界とタイル境界を再び一致させ、効率的なメモリアクセスパターンを復元できます。
パディングの設計では、複数の要因を考慮する必要があります。まず、バーストセクションサイズ(通常32バイトまたは128バイト)の倍数になるよう調整します。次に、選択したタイルサイズとの整合性を確保します。さらに、追加されるパディング量を最小限に抑えて、メモリ使用量の増加を制限します。
この手法は、nanoGPTの例で示されたように、実際のアプリケーションで劇的な効果を発揮します。語彙サイズを5257から5304(64の倍数)に調整するだけで25%の速度向上を実現したのは、まさにこのパディング戦略の成功例です。
ただし、パディングには計算コストとメモリコストのトレードオフがあります。追加要素に対する無駄な計算とメモリ使用量の増加を、アライメント最適化による性能向上が上回る場合にのみ有効です。このバランスを適切に評価し、最適なパディング戦略を選択することが、効率的なGPU最適化の鍵となるのです。
20.1 nanoGPTの劇的最適化:語彙サイズ5257→5304(64の倍数)
もちろん、torch compileやすべての行列乗算のCUDA最適化などは、私が話したような種類のことを正確に行っています。それが、より良い性能を得る方法です。
すべてのこの行列の複雑さは、このような状況に終わります。Andreのこのツイートを読んでいますが、nanoGPTに対する最も劇的な最適化は、語彙サイズを5257から5304に増やすことです。これは64の最も近い倍数であり、はるかに高いオキュパンシーを提供し、2の累乗に注意深く、47次元を語彙に追加することから25%の速度向上をもたらします。
このnanoGPTの事例は、GPU最適化における理論と実践の完璧な融合を示しています。Andrej Karpathyのツイートから引用されたこの実例は、前述したタイリング、アライメント、パディングの概念がいかに実用的な価値を持つかを具体的に実証しています。
語彙サイズ5257という数値は、自然言語処理の観点からは合理的な選択だったかもしれません。しかし、GPU最適化の視点では、この数値は複数の問題を抱えていました。まず、64の倍数ではないため、メモリアライメントが最適化されていませんでした。また、2の累乗系列から外れているため、バーストセクション境界との整合性に問題がありました。
5304への調整(5257 + 47 = 5304 = 64 × 83)により、これらの問題が一挙に解決されました。64という数値は、多くのGPUアーキテクチャにおいて重要な意味を持ちます。ワープサイズ(32)の倍数であり、典型的なバーストセクションサイズとも整合し、テンソルコアの最適な動作パターンとも一致します。
「はるかに高いオキュパンシー」という表現は、SMリソースの利用効率向上を示しています。適切にアライメントされた行列演算により、より多くのSMが効率的に動作し、アイドル時間が削減されたのです。
最も驚くべきは、「47次元を語彙に追加することから25%の速度向上」という結果です。計算内容や アルゴリズムを一切変更せずに、純粋にデータ構造の調整だけで、これほど大幅な性能向上を実現したのです。これは、GPU最適化において「注意深い2の累乗」の選択がいかに重要かを示す完璧な実証例なのです。
20.2 47次元追加による25%速度向上の驚き
これは、47次元を語彙に追加することから25%の速度向上をもたらします。たった47次元追加で25%の性能向上です。
この47次元追加による25%速度向上は、GPU最適化の反直感的な性質を極めて鮮明に示している事例です。従来のソフトウェア最適化では、データ量の増加は通常、処理時間の増加を意味します。しかし、GPUにおいては、適切なアライメントのための「無駄な」次元追加が、劇的な性能向上をもたらすのです。
この驚きの背景には、前述したGPUアーキテクチャの複雑な相互作用があります。5257から5304への47次元追加は、単なる数値変更ではなく、メモリアクセスパターン、タイルアライメント、バーストセクション効率、SM利用率など、複数の最適化要因を同時に改善したのです。
特に重要なのは、この改善が「無料」で得られたことです。追加された47次元は実際の言語モデルの能力には寄与しませんが、ハードウェア効率の大幅向上により、全体の訓練・推論速度が25%向上しました。これは、同じ計算予算でより多くの実験を実行できることを意味し、研究開発の生産性に直接的な影響を与えます。
この事例の教訓は、現代のGPU最適化では、アルゴリズムの理論的効率性だけでなく、ハードウェア特性との適合性が決定的に重要だということです。数学的に最適な解が、必ずしもハードウェア上で最適な性能を発揮するとは限りません。むしろ、ハードウェアの制約と特性を深く理解し、それに適合するようデータ構造を調整することで、予想を超える性能向上を実現できるのです。
この47次元の追加は、GPU最適化における「詳細への注意」の重要性を象徴する完璧な例であり、理論と実践の橋渡しとなる貴重な知見を提供しているのです。
20.3 2の累乗の重要性の実証
これは64の最も近い倍数であり、2の累乗に注意深く、47次元を語彙に追加することから25%の速度向上をもたらします。それは、このような種類のことがどのように発生するかです。
この実例は、GPU最適化における2の累乗の重要性を明確に実証しています。64という数値選択は偶然ではなく、GPUアーキテクチャの根本的特性に基づいた戦略的判断です。
64は複数の重要な2の累乗系列の交点に位置しています。64 = 2^6であり、ワープサイズ32(2^5)の倍数でもあります。この特性により、32スレッドのワープが語彙次元上で動作する際、完璧にアライメントされたメモリアクセスパターンを実現できます。また、多くのGPUにおけるバーストセクションサイズやキャッシュライン境界とも整合します。
「2の累乗に注意深く」という表現は、単なる数学的美学ではなく、実用的な性能要件を示しています。現代のコンピュータハードウェアは、2進数ベースの設計が根底にあるため、2の累乗に基づくデータ構造が最も効率的に処理されます。メモリアドレッシング、キャッシュ管理、並列処理の分散など、あらゆるレベルで2の累乗アライメントが最適化の鍵となります。
この原則は、nanoGPTのような単一事例だけでなく、現代のGPU最適化における普遍的な設計指針です。深層学習における隠れ層サイズ、バッチサイズ、埋め込み次元など、主要なハイパーパラメータが512、1024、2048といった2の累乗値に設定されることが多いのは、この原則に基づいています。
47次元という「小さな」調整から25%という「大きな」改善が生まれたことは、GPU最適化における2の累乗の威力を端的に示しています。これは、効率的なGPUプログラミングでは、アルゴリズムの数学的美しさよりも、ハードウェア特性との調和が優先されることを明確に示す実証例なのです。
21. 謎のパフォーマンス曲線の完全解明
21.1 第1段階:1536以下での計算強度不足(Rooflineモデル)
これで、私たちは謎に戻ることができます。GPUの詳細をすべて引きずってきたのは、皆さんがすべての性能特性を完全に理解できるようになることを願ってのことですが、ある意味でその見返りは、この図表がどのように生まれるかを説明できることです。最終的に、行列乗算性能がそれほど神秘的でも怖くもないことがわかるでしょう。
最初の部分は非常に非常に単純です。計算強度がわかります。これは、最初に指摘したrooflineです。ここまで、約1536までは、行うべき行列乗算作業が十分にありません。ロードしなければならない行列と行わなければならない非常に基本的なIOが、この点以下でボトルネックになります。
この点を過ぎると、計算ユニットをサポートするのに十分なメモリ帯域幅がありません。理論的には、右側のここでは、上部エンベロープを描くと、これは達成可能な最大性能の種類です。
すべての計算ユニットを飽和させ、本当に素晴らしい性能を得ることが可能です。しかし、行列サイズを台無しにすると、これらの種類の本当に奇妙な場所に終わることがあり、これらのそれぞれの中で、奇妙な谷に終わることがあります。
この第1段階の説明は、謎の性能曲線解明における基礎段階を明確に示しています。1536という具体的な閾値は、計算強度とメモリ帯域幅のバランスが変化する転換点です。
1536サイズ以下の領域では、行列乗算に必要な実際の計算量が相対的に少なく、GPUの強力な計算ユニットを十分に活用できません。代わりに、行列データの読み込み、メモリからのデータ転送、カーネル起動のオーバーヘッドなど、「非常に基本的なIO」が全体性能を支配します。これは典型的なメモリバウンド状態であり、rooflineモデルの左側(対角線部分)に対応します。
この段階では、いくら強力な計算ユニットがあっても、それらに供給するデータが不足しているため、計算能力を活かすことができません。200-300クロックサイクルのグローバルメモリアクセス時間に対して、実際の行列乗算計算時間が相対的に短すぎるのです。
1536を超える領域では、状況が逆転します。十分な計算作業量により、計算ユニットをフル活用できる可能性が生まれますが、今度はメモリ帯域幅が制約となります。「計算ユニットをサポートするのに十分なメモリ帯域幅がない」状況では、計算ユニットがデータ待機でアイドル状態になります。
理論的な「上部エンベロープ」は、すべての制約が理想的に調整された場合の最大性能を示しています。しかし実際には、「行列サイズを台無しにする」ことで、この理想から大きく外れた「奇妙な谷」に陥る可能性があるのです。これらの谷こそが、次に説明する複雑な最適化要因の相互作用によって生まれる現象なのです。
21.2 第2段階:タイリング・アライメント問題
これらのそれぞれの中で何が起こっているかについて少し考えてみましょう。
最初のもの、この最初の線はタイリングアライメントの問題です。ここの倍数を見ると、行列サイズの割り切れる度合いに基づいて、これらの線のそれぞれを色分けしました。これは、それが割り切れるサイズです。
32で割り切れる場合、これらの紫の点の上のここで、良い状態にあります。16で割り切れる場合、実際にはまだここの上にあります。2つの色があります。そして緑の場合、k=8で、ここの上にあります。オレンジの場合、k=2です。そしてk=1の場合、ここまで下にあります。どの数でも割り切れない場合、プライム次元を選ばないでください。行列乗算で非常に良いスループットを得ることはできません。
これの大部分は、k=2とk=1に到達すると、基本的にバースト読み取りとうまく整列した方法でタイルを読み取ることができなくなる状況を強制することになるからです。それは深刻な問題につながるでしょう。
このタイリング・アライメント問題の分析は、性能曲線の波状パターンの主要因を明確にしています。行列サイズの約数特性が、直接的に性能階層を決定しているのです。
紫色の点(k=32で割り切れる)が最高性能グループを形成しているのは、32という数値がGPUアーキテクチャにとって特別な意味を持つからです。32はワープサイズであり、多くの最適化されたタイルサイズの公約数でもあります。このサイズで割り切れる行列は、完璧なタイルアライメントを実現でき、メモリアクセスパターンも最適化されます。
k=16やk=8の場合も、依然として比較的良好な性能を維持しています。これらの数値も重要なハードウェア境界と整合するため、合理的な最適化効果を得られます。しかし、k=32ほど完璧ではないため、若干の性能低下が発生します。
問題が深刻化するのはk=2とk=1の領域です。これらの行列サイズでは、効率的なタイル分割が困難になります。特に、「バースト読み取りとうまく整列した方法でタイルを読み取ることができなくなる」状況が発生します。前述したバーストセクションアライメント問題により、メモリアクセス効率が大幅に低下し、最悪の場合は2倍のメモリアクセス時間が必要になります。
最下位のk=1(プライム次元)では、ほとんどあらゆる最適化手法が無効化されます。タイル分割、メモリアライメント、効率的な並列化など、GPU最適化の基本要素がすべて機能しなくなるため、「行列乗算で非常に良いスループットを得ることはできません」という結果になります。
この分析により、なぜ実用的なGPUアプリケーションで2の累乗や豊富な約数を持つ次元が好まれるのかが明確になります。これは単なる慣習ではなく、ハードウェア特性に基づいた必然的な選択なのです。
21.3 約数別性能差:k=32(紫)> k=16,8(上位)> k=2(オレンジ)> k=1(最下位)
ここの倍数を見ると、行列サイズの割り切れる度合いに基づいて、これらの線のそれぞれを色分けしました。これは、それが割り切れるサイズです。
32で割り切れる場合、これらの紫の点の上のここで、良い状態にあります。16で割り切れる場合、実際にはまだここの上にあります。2つの色があります。そして緑の場合、k=8で、ここの上にあります。オレンジの場合、k=2です。そしてk=1の場合、ここまで下にあります。どの数でも割り切れない場合、プライム次元を選ばないでください。行列乗算で非常に良いスループットを得ることはできません。
この約数別性能階層は、GPU最適化における数学的美しさとハードウェア現実の関係を明確に示しています。性能差は単なる微細な調整ではなく、明確に区分された階層構造を形成しているのです。
最上位のk=32(紫色)グループは、GPUアーキテクチャとの完璧な調和を達成しています。32という数値は、ワープサイズ(32スレッド)、一般的なタイルサイズ(32の倍数)、メモリアライメント境界など、複数の重要なハードウェア特性と同期しています。この完璧な整合により、メモリアクセス効率、計算ユニット利用率、並列化効果のすべてが最適化されます。
k=16とk=8の上位グループは、依然として重要なハードウェア境界と整合するため、高い性能を維持します。これらの数値も2の累乗であり、多くの最適化手法と互換性があります。ただし、k=32ほど完璧ではないため、若干の性能ギャップが存在します。
k=2(オレンジ)では、最適化オプションが大幅に制限されます。バーストセクションアライメント、効率的なタイル分割、メモリ合体などの恩恵を十分に受けることができません。結果として、上位グループとの間に明確な性能ギャップが生じます。
最下位のk=1は、数学的には完全に有効な行列サイズですが、GPU最適化の観点では最悪の選択です。プライム数次元では、ほぼすべての並列化とメモリ最適化手法が無効化されます。タイル分割時の境界効果、非効率的なメモリアクセス、不均等な負荷分散など、あらゆる問題が同時に発生します。
この階層的性能差は、現代のGPU プログラミングにおいて、アルゴリズムの数学的正確性だけでなく、ハードウェア適合性が決定的に重要であることを実証しています。最適な約数特性を持つ次元選択は、同じ計算内容でも劇的に異なる性能を生み出すのです。
21.4 プライム次元の回避推奨
そしてk=1の場合、ここまで下にあります。どの数でも割り切れない場合、プライム次元を選ばないでください。行列乗算で非常に良いスループットを得ることはできません。
これの大部分は、k=2とk=1に到達すると、基本的にバースト読み取りとうまく整列した方法でタイルを読み取ることができなくなる状況を強制することになるからです。それは深刻な問題につながるでしょう。
このプライム次元回避の推奨は、GPU最適化における最も実用的で重要なガイドラインの一つです。数学的に完全に有効なプライム数が、ハードウェア上では最悪の性能を示すという現実は、理論と実装の間の深刻な乖離を象徴しています。
プライム次元の根本的問題は、約数が1と自分自身しか存在しないことです。この特性により、あらゆる分割ベースの最適化手法が無効化されます。タイル化では効率的な分割ができず、必然的に大きな余剰領域を持つスパースなタイルが発生します。これらのタイルを担当するSMは、ほとんど何もすることがなく、貴重な計算リソースが無駄になります。
さらに深刻なのは、メモリアクセスパターンの問題です。プライム次元では、バーストセクション境界との整合が極めて困難になります。「バースト読み取りとうまく整列した方法でタイルを読み取ることができなくなる」結果、前述した2倍メモリアクセス問題が頻発します。4倍のスループット向上機会を完全に失い、むしろ性能が大幅に低下してしまいます。
この問題は、プライム数の数学的美しさとは裏腹に、実用的な計算において致命的な欠陥となります。例えば、行列サイズが1009(プライム数)の場合、どのような分割を行っても効率的なタイリングは不可能で、メモリアクセスも非効率的になります。
「深刻な問題につながる」という警告は、単なる性能低下だけでなく、開発効率やデバッグの困難さも含んでいます。プライム次元による予期しない性能低下は、しばしば他の要因として誤認され、無駄な最適化努力を引き起こします。
この推奨の実用的価値は、nanoGPTの事例でも実証されています。プライム次元を避け、豊富な約数を持つ次元を選択することで、大幅な性能向上を簡単に実現できるのです。
22. Wave量子化現象の詳細分析
22.1 1792→1794サイズ変更での劇的性能低下の謎
それはミステリーの一部ですが、ミステリーの別の部分が残っていると思います。この内部のオレンジ線の中で、ここにズームインすると、この巨大な落下が見えます。この点からこの点まで、次元を2つ増やしただけでどうしてこれほど多くの性能を失うことができるのか疑問に思っているでしょう。
これらの数字をただ見てみましょう。これは楽しいパズルだと思います。1792から1794の4つのサイズに移行するときに、これが起こります。これを4としましょう。これは2の因数のままです。
この1792から1794への変更における劇的な性能低下は、GPU最適化の最も反直感的な現象の一つです。わずか2の増加(実際には講師は後で1794ではなく1794の4、つまり1796について言及していますが、これを4としてより分析しやすくしています)で、なぜこれほど大幅な性能低下が発生するのでしょうか。
表面的には、この変更は些細に見えます。1792も1796も2の倍数であり、一見すると似たような数値特性を持っています。実際の計算量もほぼ同じで、追加される要素はわずかです。従来のCPUベースの計算では、このような小さな変更が大幅な性能変化を引き起こすことは稀です。
しかし、GPUの複雑な最適化メカニズムでは、この「小さな」変更が複数のシステムレベルの問題を連鎖的に引き起こします。行列サイズの変更は、タイル分割パターン、SMへの作業配分、メモリアクセスパターン、さらには物理的なハードウェアリソースの割り当てにまで影響を与えるのです。
このパズルの解明は、GPU性能の理解において重要な意味を持ちます。なぜなら、こうした予期しない性能変化は、実際のアプリケーション開発において頻繁に遭遇する問題だからです。研究者や開発者が行列サイズを微調整する際に、理論的には改善されるはずの変更が、実際には劇的な性能低下を引き起こすことがあります。
この現象を理解することで、GPU最適化における「詳細への注意」がいかに重要かが明確になります。数値の表面的な類似性に惑わされず、ハードウェアレベルでの相互作用を深く理解する必要があるのです。
22.2 256×128タイルサイズでの計算:98タイル → 120タイル増加
これらの数字をただ見てみましょう。これは楽しいパズルだと思います。1792から1794の4つのサイズに移行するときに、これが起こります。これを4としましょう。これは2の因数のままです。なぜそれが起こるのでしょうか?
256×128のタイルサイズを使用しているとしましょう。これは非常に自然なサイズです。楽しい事実として、これらのGPUの行列乗算ユニットは、自然に約128のサイズの行列で動作しています。そのため、256×128は非常に良いタイルサイズです。
つまり、いくつのタイルがありますか?7×14のタイルがあります。行列の次元をタイルのサイズで割っているからです。それは合計98の異なるタイルです。
これを1つ増やすと、各座標を切り上げなければならなくなります。そうすると、多くのタイルが増え、120個のタイルになります。タイルの数をかなり増やしました。
この計算は、Wave量子化現象の数値的基盤を明確に示しています。256×128というタイルサイズの選択は、単なる恣意的な決定ではなく、GPUハードウェアの物理的特性に基づいた最適化された選択です。
「楽しい事実」として言及されているように、GPUの行列乗算ユニット(テンソルコア)は、約128サイズの行列で最適に動作するよう設計されています。これは、ハードウェアレベルでの並列処理ユニットの構成と直接関連しています。256×128タイルは、この物理的制約を活用しつつ、共有メモリの制限内で最大限のデータを処理できるよう設計された「非常に良いタイルサイズ」なのです。
1792サイズの場合、各次元でのタイル数は1792÷256=7と1792÷128=14となり、7×14=98タイルで完全に分割されます。これは非常に効率的な分割で、境界での無駄がほとんどありません。
しかし、1796サイズ(1792+4)に変更すると、状況が劇的に変化します。各次元でのタイル数は、1796÷256=7.015...≈8(切り上げ)と1796÷128=14.03...≈15(切り上げ)となり、8×15=120タイルが必要になります。
この98から120への22タイル増加(約22%の増加)は、単なる数値変更ではありません。追加されたタイルの多くは、実際のデータをほとんど含まない「スパース」な状態となり、計算効率を大幅に低下させます。さらに重要なのは、この増加がGPUの物理的リソース制限と相互作用して、次に説明するSM割り当ての問題を引き起こすことです。
22.3 A100の108 SM制限:98タイル(完全利用)vs 120タイル(12タイル低利用待機)
さて、そこで、利用率だけでなく、さらに悪いことに、A100は108のSMを持っています。GPU実行モデルに戻ると、SMは並列に実行でき、それらは実行ユニットの一種です。98のSMがある場合、それらはすべて行って実行されます。すべてのSMが実行されており、優れた利用率を得ています。
120のタイルに移行すると、SMよりもタイルが多くなりました。108のタイルが実行され、非常に低い利用率で残りの12を実行するために戻って、それらの完了を待つことになります。これはwave quantizationと呼ばれるものです。
理想的には、タイルサイズはSMの数よりもはるかに大きいか、SMを超えてこの量子化エラーを引き起こさない場所にあることが望ましいです。
このA100のSM制限による問題は、Wave量子化現象の核心を成しています。108という物理的なSM数は、A100 GPUのハードウェア制約であり、変更することはできません。この制約が、タイル数の微小な変化を劇的な性能問題に変換するのです。
98タイルの場合、すべてのタイルが108個の利用可能なSMに並列配布されます。各SMが1個のタイルを処理し、10個のSMは待機状態になりますが、全体として効率的な並列実行が実現されます。重要なのは、すべてのタイルが同時に処理されるため、同期待機や段階的実行による遅延が発生しないことです。
しかし、120タイルでは状況が根本的に変化します。第1段階で108個のタイルが並列実行されますが、この段階が完了した後、残り12個のタイルを処理するための第2段階が必要になります。この第2段階では、12個のSMのみが動作し、残り96個のSMは完全にアイドル状態となります。
この段階的実行により、「非常に低い利用率で残りの12を実行する」状況が発生します。第2段階では、GPU全体の計算能力の約11%(12/108)しか活用されません。さらに深刻なのは、第1段階で完了した108個のSMが、第2段階の完了まで待機しなければならないことです。
「Wave quantization」という名称は、この現象が波のような段階的パターンを示すことに由来します。理想的な連続的利用率ではなく、高利用率→低利用率→待機という量子化されたパターンが発生するのです。
この現象の教訓は、効率的なGPUプログラミングでは、タイル数がSM数との関係を慎重に設計する必要があることです。理想的には、タイル数がSM数を大幅に上回る(多数の小さなタイル)か、適切に約数関係を持つよう設計することで、このような量子化エラーを回避できるのです。
22.4 利用率パターン:高利用→急降下→低利用完了
120のタイルに移行すると、SMよりもタイルが多くなりました。108のタイルが実行され、非常に低い利用率で残りの12を実行するために戻って、それらの完了を待つことになります。
利用率を見ると、しばらくの間は良い利用率を得て、それから崖から落ちて、それからあなたの仕事を終えるでしょう。これはwave quantizationと呼ばれるものです。
この利用率パターンの説明は、Wave量子化現象の実際の動作メカニズムを時系列で明確に示しています。GPU全体の利用率が時間とともにどのように変化するかを理解することで、なぜこの現象が深刻な性能問題となるのかが分かります。
第1段階(高利用率期間)では、108個のSMすべてが同時に動作し、各SMが1つのタイルを処理します。この期間中、GPU全体の利用率は約90%(98タイル分の有効作業)から100%近くに達します。計算ユニット、メモリ帯域幅、その他のリソースがすべて効率的に活用され、理想的な並列実行状態が実現されます。
しかし、108個のタイル処理が完了した瞬間、「崖から落ちる」ような急激な利用率低下が発生します。この移行は段階的ではなく、瞬間的です。108個のSMが同時に作業を完了し、次の瞬間には12個のSMのみが残り12個のタイルを処理する状態に移行します。
第2段階(低利用率期間)では、利用率が約11%(12/108)まで急落します。96個のSM(約89%のリソース)が完全にアイドル状態となり、何もすることなく待機します。この期間は比較的短いものの、GPU全体の性能に深刻な影響を与えます。なぜなら、全体の実行時間は最も遅い段階によって決定されるからです。
「あなたの仕事を終える」最終段階では、残り12個のタイルが完了するまで、すべてのSMが待機状態になります。この同期待機により、実質的にGPU全体の実行時間が延長されます。
このパターンが「wave quantization」と呼ばれる理由は、利用率が波のように高低を繰り返し、かつ量子化された(離散的な)ステップで変化するからです。連続的で安定した利用率ではなく、急激な変化を伴う不連続なパターンが、全体性能を大幅に悪化させるのです。
23. GPU最適化技術の総括
23.1 メモリアクセス削減手法:合体、融合、共有メモリ移動、タイリング
トリックは何でしたか?重要なアイデアです。まず、メモリアクセスの量を減らす必要があります。それを行う多くの方法があります。コアレッシングを行うことができます。無料で取得しているリードを再利用できるように、アクセスをグループ化することです。
融合を行うことができます。複数の操作を一緒に融合して、不必要な読み書きを避けることができます。メモリを共有メモリに移動できます。読み書きを行う場合でも、それらがはるかに高速なメモリからのものになります。これがあなたが行うことができるタイリングトリックです。
このメモリアクセス削減手法の総括は、GPU最適化における4つの主要戦略を体系的に整理しています。これらの手法はすべて、根本的な問題である「メモリ帯域幅がボトルネック」という現実に対する異なるアプローチを提供します。
コアレッシング(合体)は、ハードウェアの物理的特性を活用する手法です。DRAMのバーストモードにより、1回のメモリアクセスで複数の連続データが「無料で」取得されます。この特性を活用してアクセスパターンを最適化することで、同じデータ量を取得するのに必要なメモリアクセス回数を4分の1に削減できます。32個のスレッドからなるワープが連続アドレスにアクセスすることで、個別アクセスではなく効率的な一括取得が実現されます。
融合は、演算レベルでの最適化です。sin²x + cos²xの例で示されたように、本来5つの個別カーネルが必要な計算を1つのカーネルに統合することで、中間結果の無駄なメモリ往復を完全に排除します。torch compileのような自動化ツールにより、この最適化は開発者にとってアクセシブルになっています。
共有メモリ移動は、メモリ階層の活用です。200-300クロックサイクルのグローバルメモリアクセスを、20クロックサイクルの共有メモリアクセスに置き換えることで、約10倍の高速化を実現します。データを一度共有メモリに読み込めば、その後の全てのアクセスが高速になります。
タイリングは、これらの手法を統合した包括的戦略です。データを小さなタイルに分割し、各タイルを共有メモリに読み込んで集中的に処理することで、グローバルメモリアクセスを大幅に削減します。理論的にはt倍(タイルサイズ)の削減効果が得られ、大きなタイルほど効果が高くなります。
これら4つの手法は独立した技術ではなく、相互に補完し合う統合されたアプローチとして機能します。効率的なGPUプログラムでは、通常これらすべての手法が組み合わせて使用されるのです。
23.2 リソーストレードオフ:メモリ↔計算(再計算)、メモリ↔精度(量子化)
最後に、他のリソースを持っているものとメモリを交換することができます。再計算になる計算のためにそれを交換するか、量子化になる数値精度や安定性のためにそれを交換することができます。あなたが持っているリソース間の戦略的なトレードオフの多くの種類のバッグオブトリックがあります。
このリソーストレードオフの概念は、現代GPU最適化における最も重要な戦略的思考の一つです。GPUアーキテクチャの不均衡な発展(計算能力の爆発的成長 vs メモリ帯域幅の緩やかな成長)により、異なるリソース間での意図的な交換が極めて有効になっています。
メモリ↔計算(再計算)のトレードオフは、3重シグモイド関数の例で実証されました。従来は中間活性化をグローバルメモリに保存していましたが、代わりにこれらの値を必要時に再計算することで、8回から5回へのメモリアクセス削減(37.5%改善)を実現しました。数クロックサイクルの追加計算により、200-300クロックサイクルのメモリアクセスを回避できるため、このトレードオフは圧倒的に有利です。
重要なのは、この手法が現代のGPU環境で特に効果的な理由です。計算能力が余剰となっている一方でメモリ帯域幅が不足している状況では、「計算のためにメモリを交換する」ことで全体性能が向上します。アイドル状態の計算ユニットを有効活用し、ボトルネックとなっているメモリアクセスを削減できるのです。
メモリ↔精度(量子化)のトレードオフは、FP32からFP16への移行例で示されました。数値精度を部分的に犠牲にすることで、メモリ使用量とメモリ帯域幅を半減でき、実質的に「無料でメモリ帯域幅を2倍」にできます。混合精度アプローチにより、重要な計算では高精度を維持しつつ、大部分の処理で低精度の恩恵を受けることが可能です。
これらのトレードオフが「バッグオブトリック」と表現されているのは、状況に応じて適切な手法を選択する必要があるからです。ワークロードの性質、ハードウェアの制約、精度要件などを総合的に評価し、最適なリソース配分を決定することが、効率的なGPU最適化の核心なのです。
23.3 メモリがGPU性能の鍵である理由の総合理解
メモリについて本当に注意深く考える必要があるだけです。これがGPUの動作方法を考える上でのある意味での鍵であることです。
そして最後に、一種のデータ移動を最適化したいと思います。特定の量の計算を行わなければならない場合、物事を最適化する方法は、高帯域幅メモリまたはグローバルメモリからできるだけ多くの移動を避け、すべてを非常に非常に高速な共有メモリに配置し、それが良いパフォーマンスにつながることです。
このメモリ中心思考の重要性は、講義全体を通じて蓄積された洞察の集大成です。GPUコンポーネントのスケーリング比較で示されたように、計算性能(灰色線)が1〜100,000倍という驚異的成長を遂げた一方、メモリ帯域幅(緑線)は相対的に緩やかな成長にとどまっています。この根本的な不均衡こそが、メモリがGPU性能の決定的要因となっている理由です。
現代のH100のような最新GPUでは、計算ユニットは膨大な処理能力を持っていますが、それを支えるメモリ供給が追いついていません。結果として、多くの実際のワークロードで計算ユニットがデータ待機でアイドル状態になり、高価で高性能なハードウェアが十分に活用されない状況が発生しています。この現実により、「FLOPs削減よりもメモリ移動効率化」が優先されるようになったのです。
「データ移動を最適化したい」という表現は、この新しいパラダイムを端的に示しています。従来のCPU最適化では演算数の削減が主要目標でしたが、GPU最適化では演算数が同じでもメモリ移動パターンを改善することで劇的な性能向上を実現できます。nanoGPTの事例(47次元追加で25%向上)は、計算内容を一切変更せずに、純粋にメモリアクセス最適化だけで大幅改善を達成した完璧な実例です。
最適化戦略の核心は明確です:「高帯域幅メモリまたはグローバルメモリからできるだけ多くの移動を避け、すべてを非常に非常に高速な共有メモリに配置する」こと。200-300クロックサイクルの遅いアクセスを20クロックサイクルの高速アクセスに置き換えることで、約10倍の性能向上を実現できます。
この理解は、後に説明されるFlash Attentionのような先進的アルゴリズムの基盤でもあります。メモリアクセスを準二次化することで、計算量は二次のままでも全体性能を劇的に改善できるのです。メモリ中心思考こそが、現代GPU最適化における最も重要な設計原則なのです。
24.1 概要:CUDAカーネル最適化による劇的な注意機構高速化
さて、すべてをまとめます。私が教えたトリックが、GPUについてのランダムな切り離された事実ではないようにします。それらは、標準的な性能最適化ツールキットの一部であり、Flash AttentionとFlash Attention 2が、うまくいけば、現代の高性能トランスフォーマーの基盤の一つである、それがすべてどのように組み合わさるかを教えてくれるでしょう。
Flash Attentionは、注意機構を劇的に加速することを知っています。皆さんのほとんどは、それがいくつかのCUDAカーネルマジックを通じて行われることを知っているでしょうが、おそらくすべての詳細は知らないでしょう。
論文が言うところでは、さて、起こっている一つの部分があります。最適化されていないPyTorchトランスフォーマー実装で注意を行い、カーネルを融合し、いくつかのことを行うと、大幅な大幅な速度向上を得ることができます。
このFlash Attentionの概要説明は、講義全体の学習内容を統合する重要な転換点を示しています。これまで説明してきたGPU最適化技術が、「ランダムな切り離された事実」ではなく、実際の革新的アルゴリズムにおいて体系的に活用される実例を提供するのです。
Flash Attentionの革新性は、単純に「CUDAカーネルマジック」として片付けられがちですが、実際にはこれまで学習した複数の最適化技術の精密な組み合わせによって実現されています。注意機構という、現代のTransformerアーキテクチャにおいて最も計算集約的で重要な構成要素を、根本的に再設計することで劇的な性能向上を達成したのです。
「現代の高性能トランスフォーマーの基盤の一つ」という位置づけは、この技術の実用的重要性を示しています。ChatGPTやGPT-4のような大規模言語モデルの実用化において、Flash Attentionのような最適化技術が決定的な役割を果たしています。理論的に美しいアルゴリズムも、効率的な実装なしには実用的価値を持ちません。
論文の主張する「大幅な大幅な速度向上」は、単なる数値的改善以上の意味を持ちます。この高速化により、従来は計算コストの制約で不可能だった長いコンテキスト処理や、大規模モデルの実用的展開が可能になりました。Flash Attentionは、理論的限界を技術的革新で突破した好例なのです。
重要なのは、この成功が「カーネル融合といくつかのこと」という表現で示されるように、これまで学習した基本技術の応用によって実現されていることです。革新的な成果も、基礎的な最適化原理の深い理解と巧妙な組み合わせから生まれるのです。
24.2 論文の主張:タイリング+再計算による準二次HBMアクセス実現
論文から、彼らは2つの確立された技術、タイリングと再計算を適用して、準二次HBMアクセスで正確な注意を計算するという技術的課題を克服すると言っています。準二次計算ではなく、なぜなら一般的にそれを行うことはできないからです。注意を行う必要がありますが、HBMまたは高帯域幅へのアクセスが準二次になります。
メモリがボトルネックである場合、そのメモリアクセスを二次的にしたくありません。少なくとも計算よりもメモリでその二次コストを支払うことができます。
この論文の核心的主張は、Flash Attentionの技術的革新が既存の確立された最適化手法の巧妙な組み合わせであることを明確にしています。「2つの確立された技術」という表現は、完全に新しい理論的突破ではなく、既知の最適化原理の創造的応用であることを示しています。
タイリングと再計算という選択は、これまでの講義内容を考慮すると非常に理にかなっています。タイリングにより、大きな注意行列を小さな管理可能なブロックに分割し、各ブロックを高速な共有メモリで処理できます。再計算により、中間結果の保存によるメモリ使用量を削減し、必要時にオンザフライで値を再生成します。
「準二次HBMアクセス」という表現が重要な区別を示しています。注意機構の本質的な計算複雑度O(n²)は変更できません。n×nの全ての要素間相互作用を計算する必要があるためです。しかし、メモリアクセスパターンは最適化可能です。従来の実装では、HBM(High Bandwidth Memory)への読み書きもO(n²)でしたが、Flash AttentionではこれをO(n²/B)(Bはブロックサイズ)に削減しています。
この区別が決定的に重要な理由は、現代GPUにおけるボトルネックの性質にあります。前述したように、計算能力は爆発的に成長したがメモリ帯域幅は相対的に制限されています。「メモリがボトルネックである場合、そのメモリアクセスを二次的にしたくありません」という表現は、この現実を端的に示しています。
「少なくとも計算よりもメモリでその二次コストを支払うことができます」という戦略は、リソーストレードオフの完璧な実例です。余剰となった計算能力を活用して、制約となっているメモリアクセスを削減する。これこそが、現代GPU最適化における最も重要な設計思想なのです。
24.3 メモリボトルネック解決:二次計算コストは許容、二次メモリコストは回避
準二次計算ではなく、なぜなら一般的にそれを行うことはできないからです。注意を行う必要がありますが、HBMまたは高帯域幅へのアクセスが準二次になります。
メモリがボトルネックである場合、そのメモリアクセスを二次的にしたくありません。少なくとも計算よりもメモリでその二次コストを支払うことができます。
このメモリボトルネック解決の戦略は、現代GPU最適化における最も重要なパラダイムシフトを表現しています。従来のアルゴリズム設計では、計算複雑度の削減が最優先でしたが、Flash Attentionは全く異なるアプローチを採用しています。
「準二次計算ではなく、なぜなら一般的にそれを行うことはできない」という制約認識が出発点です。注意機構では、クエリとキーの間のすべてのペアワイズ相互作用を計算する必要があり、これは本質的にO(n²)の計算量を要求します。この数学的必然性は回避できません。
しかし、「HBMまたは高帯域幅へのアクセスが準二次になる」という部分に革新があります。同じ計算結果を得るために必要なメモリアクセスパターンは最適化可能です。従来の実装では中間結果をすべてHBMに保存していましたが、Flash Attentionはタイリングと再計算により、これらの中間結果を高速な共有メモリ内で処理し、HBMアクセスを大幅に削減します。
「メモリがボトルネックである場合、そのメモリアクセスを二次的にしたくありません」という原則は、講義全体で強調されてきたメモリ中心思考の完璧な実例です。計算能力が1〜100,000倍成長した一方でメモリ帯域幅は相対的に制限されている現実において、メモリアクセスこそが性能を決定する要因となっています。
「少なくとも計算よりもメモリでその二次コストを支払うことができます」という表現は、リソーストレードオフの戦略的判断を示しています。現代のGPUでは計算ユニットが余剰となることが多いため、追加的な計算コストは相対的に安価です。一方、200-300クロックサイクルを要するHBMアクセスは極めて高価なため、計算の増加と引き換えにメモリアクセスを削減することで、全体性能が大幅に向上するのです。
この戦略は、Flash Attentionだけでなく、現代の効率的なGPUアルゴリズム設計における普遍的な原則となっています。メモリ効率を最優先とし、計算効率は二次的に考慮するという、従来とは逆転した優先順位が新しい標準となっているのです。
25. Flash Attentionの技術詳細
25.1 標準的な注意機構:3つの行列乗算とSoftMaxの組み合わせ
本当に簡単な復習のために、この時点で、皆さんは多くの授業で注意を何度も何度も実装したことがあります。3つの異なる行列乗算があります。K、Q、Vがあり、間にsoftmaxがあります。
行列乗算はかなり単純で、タイリングで行うことができます。そのような例を見せました。注意について何が異なるのでしょうか?本当に厄介なビットになるsoftmaxのものがあります。そして、softmaxを扱うことができれば、私が話していた行列乗算のすべてのものが作用するでしょう。
この標準的注意機構の構造説明は、Flash Attentionの技術的挑戦を理解するための重要な前提を提供しています。多くの深層学習コースで実装経験があることを前提として、基本構造を簡潔に復習しています。
3つの行列乗算(K、Q、V)は、注意機構の核心的な構成要素です。クエリ行列Qとキー行列Kの乗算により親和性スコアを計算し、SoftMax正規化を経て、値行列Vとの乗算で最終的な出力を生成します。この一連の処理がTransformerの自己注意機構の数学的基盤となっています。
「行列乗算はかなり単純で、タイリングで行うことができます」という指摘は、これまでの講義内容との連続性を示しています。K×Q、SoftMax結果×Vという行列演算は、前述したタイリング最適化技術を直接適用できます。256×128のようなタイルサイズで分割し、共有メモリを活用し、メモリアクセスを最適化することで、効率的な実装が可能です。
しかし、「本当に厄介なビット」としてSoftMaxが特別に言及されているのは、この演算が従来の行列乗算最適化では対処できない根本的な問題を抱えているからです。SoftMaxは要素別演算ではなく、行全体にわたるグローバル演算(全要素の合計による正規化)を必要とします。
この問題の核心は、タイリング戦略との根本的な非互換性にあります。タイリングでは、小さなブロック内で計算を完結させることで効率化を図りますが、SoftMaxは行全体の情報(最大値の特定、指数和の計算)を必要とするため、タイル境界を超えた処理が不可欠になります。
「softmaxを扱うことができれば、私が話していた行列乗算のすべてのものが作用するでしょう」という表現は、問題の所在を明確にしています。SoftMax以外の全ての要素は既知の最適化技術で解決可能であり、SoftMaxこそがFlash Attentionにおける技術的ブレークスルーの焦点なのです。
25.2 行列乗算部分:既習のタイリング技術で対応可能
行列乗算は、私が言ったように、まさに私が教えたものです。Flash Attention論文の図1を見ると、これは本当にただの単純なタイル化された行列乗算です。K行列、Q行列が小さなブロックに切り分けられ、小さなブロックがSRAMにコピーされ、それらが乗算され、そしてそれらが、HBMに送られて、softmaxを行う場所で、そしてVと乗算します。
これは、KQV行列乗算に関して、すべて本当に単純です。しかし、今度はsoftmaxについて考えなければなりません。softmaxで何が起こっているのでしょうか?
この行列乗算部分の説明は、Flash Attentionの技術的複雑さがどこに集中しているかを明確に示しています。既習のタイリング技術が直接適用可能な部分と、新たな技術的革新が必要な部分を明確に区別しています。
Flash Attention論文の図1の詳細説明は、これまでの講義で学習したタイリング概念の実践的応用を示しています。K行列とQ行列を「小さなブロック」に分割することで、前述した256×128タイルのような管理可能なサイズに調整しています。これにより、共有メモリ(SRAM)の容量制限内で効率的な処理が可能になります。
「小さなブロックがSRAMにコピーされ、それらが乗算され」という処理フローは、理想的なタイリング実行パターンそのものです。高速なSRAM(共有メモリ)にデータを一括読み込みし、そこで集中的に計算を行い、グローバルメモリ(HBM)アクセスを最小化しています。これは、「グローバルメモリから共有メモリに断片を読み込み、共有メモリで大量の計算を行う」という前述の理想的実行モデルの完璧な実装例です。
重要なのは、この段階では「すべて本当に単純」だということです。K×Q行列乗算は、lectures で詳しく説明されたタイリング数学(t分の1削減効果)、メモリ合体(バーストセクション最適化)、共有メモリ活用などの確立された最適化技術を直接適用できます。新たな理論的ブレークスルーは不要で、既存の技術を適切に実装するだけで高い効率を達成できます。
この段階で「HBMに送られて、softmaxを行う場所で、そしてVと乗算します」という処理フローの説明により、次の重要な転換点が明確になります。V行列との乗算も既知の行列乗算技術で対応可能ですが、「softmaxで何が起こっているのでしょうか?」という問いかけこそが、Flash Attentionの真の技術的挑戦の始まりなのです。
この明確な区分により、Flash Attentionの革新性が、全く新しい手法の発明ではなく、既知の最適化技術と新しいSoftMax処理手法の巧妙な組み合わせにあることが分かります。
25.3 SoftMaxの課題:グローバル演算(行全体の合計が必要)のタイル内処理
softmaxで何が起こっているのでしょうか?ここでの重要なことはsoftmaxです。申し訳ありません、一歩戻ります。softmaxの問題は何ですか?それはグローバル演算です。注意におけるsoftmaxは行ごとに動作します。
正規化項のsoftmaxを計算するために、行全体を合計する必要があります。それは非常に問題があります。タイルがある場合、理想的には、タイル内ですべてを行いたいです。大きな行列に書き戻したくありません。タイル内で計算できるsoftmaxが必要で、オンラインでタイルごとに計算できるsoftmaxが必要です。
このSoftMaxの課題説明は、Flash Attentionにおける技術的核心部分の問題設定を明確にしています。SoftMaxがなぜ「非常に問題がある」のかを、タイリング最適化の文脈で具体的に説明しています。
「グローバル演算」としての性質こそが根本的問題です。通常の要素別演算(ReLU、tanh等)では、各要素を独立して処理できるため、タイル内で完結した処理が可能です。しかし、SoftMaxでは各要素の計算に行全体の情報が必要になります。具体的には、数値安定性のための最大値と、正規化のための指数和を行全体から計算する必要があります。
「行全体を合計する必要があります」という要求は、タイリング戦略と根本的に矛盾します。例えば、長さ10,000の行を128要素のタイルに分割した場合、約78個のタイルが必要になります。従来の実装では、すべてのタイルを処理してからグローバルに合計計算を行い、再度すべてのタイルを正規化する必要がありました。
この問題がタイリング最適化にとって「非常に問題がある」理由は、効率化の根本原理を破綻させるからです。「理想的には、タイル内ですべてを行いたい」という目標に対して、SoftMaxは必然的にタイル境界を超えた処理を要求します。「大きな行列に書き戻したくありません」という願望とは正反対に、中間結果をグローバルメモリに保存し、全タイル処理後に再度読み込む必要が生じます。
この制約により、SoftMaxを含む注意機構では、タイリングによるメモリアクセス削減効果が相殺されてしまいます。200-300クロックサイクルの高コストなHBMアクセスが大量に発生し、共有メモリの恩恵を活用できません。
「タイル内で計算できるsoftmaxが必要で、オンラインでタイルごとに計算できるsoftmaxが必要です」という要求設定は、Flash Attentionの技術的ブレークスルーの方向性を示しています。この要求を満たすアルゴリズムこそが、次に説明されるオンラインSoftMaxなのです。
26. オンラインSoftMaxアルゴリズム
26.1 バッチ版SoftMax:全x₁〜xₙ要求の制約
ここでの重要なことは、オンラインsoftmaxを使用することです。それは何ですか?値のストリームがある場合、通常のsoftmaxのバッチ版では、すべてのx1からxnを取り、それらを指数化し、合計し、除算します。それは通常のsoftmaxで行うことです。
そして、おそらく最大値を計算し、数値的に安定するためにそれを引くでしょう。これが左側の標準的な数値的に安定したsoftmaxです。
このバッチ版SoftMaxの制約説明は、従来手法の根本的限界を明確にしています。「全x₁〜xₙ要求」という制約こそが、タイリング最適化との互換性を阻害する主要因なのです。
標準的なSoftMax実装では、出力の任意の要素を計算するために、入力シーケンス全体の事前取得が必要です。これは数学的定義に由来する必然的な要求です:softmax(xi) = exp(xi - max(x)) / Σexp(xj - max(x))。この式を計算するには、最大値max(x)と正規化項Σexp(xj)の両方について、全要素x₁〜xₙの情報が必要になります。
「すべてのx1からxnを取り、それらを指数化し、合計し、除算します」という処理フローは、3段階の全要素処理を示しています:①全要素の指数化、②指数値の合計、③各要素の正規化。各段階で全データへのアクセスが必要なため、タイル化された部分処理では対応できません。
数値安定性のための最大値減算(「最大値を計算し、数値的に安定するためにそれを引く」)は、制約をさらに複雑化します。max(x)の計算にも全要素の事前取得が必要で、その後の指数化と正規化でも再度全要素にアクセスする必要があります。
この制約により、従来のSoftMax実装では効率的なタイリングが不可能でした。10,000要素の行を128要素のタイルに分割しても、各タイルの処理前に全10,000要素をメモリから読み込む必要があり、タイリングによるメモリアクセス削減効果が完全に相殺されてしまいます。
「左側の標準的な数値的に安定したsoftmax」という言及は、この制約が数学的正確性と数値安定性の要求から生じていることを示しています。単純に制約を無視することはできず、同等の数学的結果を保証しつつ、異なる計算手順を設計する必要があります。この課題に対する解決策こそが、次に説明されるオンラインSoftMaxアルゴリズムなのです。
26.2 Mikallof & Gimmelstein (2018)のオンライン手法
オンラインsoftmax、私はMikallofとGimmelstein 2018からこれを取りました。あなたは、テレスコーピング和のような引数を介して、基本的に現在実行中の正規化項と、e to the x i minus max of x k の現在実行中のtop項を引き出すことができるということを理解することができます。
つまり、あなたが行おうとしていることは、私の現在の反復であるx1からx of jで見た現在の最大値を維持することです。そして、私の最大値が更新された場合のこの修正項も維持します。これは基本的に私の最大値を修正し、それから私のここの新しい項を追加します。
このMikallof & Gimmelstein (2018)のオンライン手法説明は、Flash Attentionの技術的ブレークスルーの数学的基盤を示しています。2018年という比較的最近の研究成果を活用していることで、Flash Attentionが最新の数学的洞察を実装レベルで応用した例であることが分かります。
「テレスコーピング和のような引数」という数学的アプローチは、従来の全要素同時処理を段階的処理に変換する鍵となる技術です。テレスコーピング和(望遠鏡状和)では、連続する項の差分を利用することで、複雑な和を簡潔な形式に変換できます。この原理をSoftMax計算に応用することで、全要素の事前取得なしに正確な結果を得ることが可能になります。
「現在実行中の正規化項と、e to the x i minus max of x k の現在実行中のtop項を引き出す」という処理は、オンライン計算の核心です。従来のバッチ処理では最終的な正規化項を計算してから各要素を処理していましたが、オンライン手法では部分的な正規化項を段階的に更新していきます。
「x1からx of jで見た現在の最大値を維持する」という状態管理が重要です。新しい要素xj+1を処理する際、これまでの最大値mjと比較し、必要に応じて更新します。この running maximum の概念により、数値安定性を保ちながら段階的処理が可能になります。
「私の最大値が更新された場合のこの修正項も維持します」という補正メカニズムは、アルゴリズムの精密さを示しています。最大値の更新は既に計算済みの指数値すべてに影響するため、適切な補正が必要です。この補正により、段階的計算でも最終結果がバッチ処理と数学的に等価になることが保証されます。
この手法の革新性は、数学的等価性を保ちながら計算順序を根本的に変更したことにあります。全要素の同時アクセスから、ストリーミング処理への転換により、タイリング最適化との互換性を実現したのです。
26.3 実行最大値とフラスティック計算による段階的処理
つまり、あなたが行おうとしていることは、私の現在の反復であるx1からx of jで見た現在の最大値を維持することです。そして、私の最大値が更新された場合のこの修正項も維持します。これは基本的に私の最大値を修正し、それから私のここの新しい項を追加します。
そのため、このd of jは、この式のトップ項である項目2をオンラインで追跡し、最後に正規化器も計算でき、それから正規化されたy of iを取得できます。このd of vは、私が必要とする正規化項自体の一種です。
この実行最大値(running maximum)とプラスティック計算(フラスティック計算)の説明は、オンラインSoftMaxアルゴリズムの動的な状態管理メカニズムを詳細に示しています。
実行最大値の維持は、数値安定性を保証する重要な要素です。各反復において「x1からx of jで見た現在の最大値」を追跡することで、新しい要素xj+1が到着した際に即座に比較・更新できます。従来のバッチ処理では全要素を事前に取得して最大値を計算していましたが、オンライン手法では段階的に最大値を更新していきます。
「私の最大値が更新された場合のこの修正項も維持します」という補正メカニズムは、アルゴリズムの数学的正確性を保証します。最大値が更新されると、既に計算済みの指数値exp(xi - old_max)をすべてexp(xi - new_max)に再調整する必要があります。この補正により、段階的計算でも最終結果が数学的に等価になります。
「d of j」という変数は、オンライン計算における重要な状態変数です。これは「この式のトップ項である項目2をオンラインで追跡」する役割を果たし、部分的な指数和を段階的に蓄積します。従来の全要素同時計算とは異なり、各要素の寄与を逐次的に加算していきます。
「最後に正規化器も計算でき、それから正規化されたy of iを取得できます」という処理フローは、オンライン手法の完全性を示しています。全要素の処理完了時点で、正規化項(d of v)が完成し、各要素の最終的なSoftMax値を計算できます。
「このd of vは、私が必要とする正規化項自体の一種です」という説明は、アルゴリズムの効率性を示しています。従来の手法では正規化項を別途計算する必要がありましたが、オンライン手法では段階的処理の副産物として正規化項が自動的に構築されます。
この段階的処理の革新性は、メモリアクセスパターンの根本的変更にあります。全要素の事前取得から、ストリーミング処理への転換により、タイリング最適化との完全な互換性を実現したのです。
26.4 ストリーム処理による大型行列の回避
ここでの重要なことは、これをオンラインで行うことができるということです。x1からxnを事前に必要としません。x1からxnのストリームが必要なだけです。そして、それは本当に重要です。なぜなら、今、私はタイルごとにsoftmaxを計算できるからです。
最終的に、softmaxを計算するためにn²の完全行列を実体化する必要がないのです。
このストリーム処理による大型行列回避の説明は、Flash Attentionの最も重要な技術的ブレークスルーを明確に示しています。オンラインSoftMaxがもたらす根本的な変化は、メモリ使用量の劇的削減にあります。
「x1からxnを事前に必要としません。x1からxnのストリームが必要なだけです」という区別は決定的に重要です。従来のバッチ処理では、計算開始前に全データをメモリに読み込む必要がありました。一方、ストリーム処理では、データを逐次的に処理できるため、任意の時点で必要なメモリは現在処理中の要素と少数の状態変数のみです。
「今、私はタイルごとにsoftmaxを計算できる」という能力こそが、Flash Attentionの核心的価値です。128要素のタイルを処理する際、そのタイル内でSoftMaxの部分計算を完了し、グローバルな状態変数(実行最大値、部分正規化項)のみを次のタイルに引き継げます。これにより、各タイルが独立した処理ユニットとして機能し、効率的なタイリング最適化が可能になります。
最も重要な成果は「softmaxを計算するためにn²の完全行列を実体化する必要がない」ことです。従来の注意機構では、シーケンス長nに対してn×nの巨大な注意行列をメモリに保存する必要がありました。n=10,000の場合、100,000,000要素の行列が必要で、FP16でも約200MBのメモリを消費します。
ストリーム処理により、この巨大行列の実体化を完全に回避できます。各タイルペアの処理時には、そのタイルサイズ(例:128×128=16,384要素)分のメモリのみが必要で、処理完了後は即座に解放されます。これにより、メモリ使用量がO(n²)からO(タイルサイズ)に劇的に削減されます。
この変化は単なる最適化ではなく、長いシーケンス処理の実用化を可能にする根本的ブレークスルーです。従来の手法では、GPUメモリ制限により処理可能なシーケンス長が厳しく制約されていましたが、Flash Attentionにより大幅に長いコンテキストの処理が可能になりました。メモリ効率の改善が、アルゴリズムの適用範囲を根本的に拡大したのです。
27. Flash Attention 2の実装ステップ
27.1 KQ行列乗算のタイル化実行
Flash Attention 2論文を見に行くと、これは私たちがあなたに実装してもらうものになります。そのため、あなたはここのこれらのステップを実際に実行することになります。まず、KQ行列乗算があり、これはタイル化されます。
これらは小さなタイル化されたチャンクであり、それらが乗算されます。softmaxをどのように計算するつもりですか?これらの指数化された合計の実行値を維持し、それから最大項について増分的に更新し、修正し続けるつもりです。
このKQ行列乗算のタイル化実行説明は、Flash Attention 2の実装における第1段階の詳細を示しています。Assignment 2での実装体験につながる実践的な内容として位置づけられています。
「まず、KQ行列乗算があり、これはタイル化されます」という処理開始は、前述した標準的タイリング技術の直接適用です。クエリ行列Qとキー行列Kを、共有メモリに収まるサイズ(例:128×128)の小さなブロックに分割します。各タイルペアは独立して処理でき、並列化の恩恵を最大限活用できます。
「これらは小さなタイル化されたチャンクであり、それらが乗算されます」という表現は、実装の具体性を示しています。各チャンクは物理的に共有メモリに読み込まれ、20クロックサイクルの高速アクセスで行列乗算が実行されます。この段階では、従来の行列乗算最適化(メモリ合体、効率的なタイル分割等)がすべて適用可能です。
重要な転換点は「softmaxをどのように計算するつもりですか?」という問いかけです。KQ乗算の結果として得られる親和性スコアに対して、従来であれば全結果をグローバルメモリに書き戻し、後でSoftMax処理を行う必要がありました。しかし、Flash Attention 2では異なるアプローチを採用します。
「これらの指数化された合計の実行値を維持し、それから最大項について増分的に更新し、修正し続けるつもりです」という戦略は、オンラインSoftMaxの実践的応用です。各KQタイルペアの乗算結果を、即座にオンラインSoftMaxアルゴリズムに供給します。グローバルメモリへの中間結果保存を完全に回避し、タイル内でSoftMax状態変数(実行最大値、部分正規化項)を段階的に更新していきます。
この実装アプローチにより、メモリアクセス効率とタイリング最適化の両方を同時に実現できます。KQ乗算とSoftMax計算が密結合された単一の処理フローとなり、中間データの無駄な移動が排除されるのです。
27.2 指数和の段階的更新と最大値補正
softmaxをどのように計算するつもりですか?これらの指数化された合計の実行値を維持し、それから最大項について増分的に更新し、修正し続けるつもりです。そして、それを行うことで、すべてのタイルを通して進み、各タイルについて、そのタイルのタイル別計算を行うことができます。
この指数和の段階的更新と最大値補正の説明は、Flash Attention 2の動的状態管理メカニズムの実装詳細を示しています。オンラインSoftMaxアルゴリズムが実際のタイル処理においてどのように機能するかを具体的に説明しています。
「これらの指数化された合計の実行値を維持し」という処理は、各タイルから得られる部分的な親和性スコアを累積的に処理することを意味します。従来の実装では全KQ乗算結果を収集してから一括でSoftMax処理していましたが、Flash Attention 2では各タイルの結果を即座に指数化し、running sum に加算します。
「最大項について増分的に更新し、修正し続ける」という動的補正は、数値安定性維持の鍵となります。新しいタイルから得られる親和性スコアが既存の最大値を上回る場合、これまでに蓄積された指数和をすべて補正する必要があります。具体的には、exp(xi - old_max)をexp(xi - new_max) = exp(xi - old_max) × exp(old_max - new_max)に変換します。
この補正メカニズムは計算負荷の増加を伴いますが、それでも全体性能が向上する理由は、メモリアクセス削減効果が補正計算コストを大幅に上回るからです。200-300クロックサイクルのグローバルメモリアクセスを回避できることを考慮すると、数クロックサイクルの追加計算は相対的に安価です。
「すべてのタイルを通して進み、各タイルについて、そのタイルのタイル別計算を行うことができます」という処理フローは、完全に分散化された実装を示しています。各タイルが独立した処理ユニットとして機能し、グローバルな依存関係は最小限の状態変数(現在最大値、部分指数和)のみに制限されます。
この設計により、タイル間の同期待機や大規模なデータ転送が不要になります。各タイルの処理完了時に更新される状態変数は数値のみで、巨大な中間行列の保存・転送は完全に排除されます。この効率化こそが、Flash Attention 2の性能向上の核心的要因なのです。
27.3 全タイル処理後の正規化項取得
学生からの質問がありました:「その出力を計算するまでは、すべてのタイルにわたってKのような2つの乗算を計算できませんよね?そのため、すべてのタイルで倍増する必要がありますね?」
そのため、すべてのタイルを通して進むまで、その分母の合計を構築しないでしょう。それは正しいです。しかし、すべてのタイルを1回行った後、たとえば、すべてのn²タイルを行った後です。その時点で、softmaxを直接出力するのに必要なすべてのコンポーネントがあります。
その時点で、これらのタイルのすべてを通した後、正規化項であるL3またはLのLのNを構築しました。私はこの最後のタイルの共有メモリに、そのようなものがすでにあります。
この全タイル処理後の正規化項取得の説明は、学生の重要な質問に対する回答を通じて、Flash Attention 2の実行フローにおける重要な時系列を明確にしています。
学生の質問「その出力を計算するまでは、すべてのタイルにわたってKのような2つの乗算を計算できませんよね?」は、オンラインSoftMaxの根本的制約を指摘しています。SoftMaxの最終出力には正規化項(分母)が必要で、これは全タイルの処理完了まで確定しません。この理解は正確で、オンライン処理の本質的な特性です。
「すべてのタイルを通して進むまで、その分母の合計を構築しないでしょう。それは正しいです」という確認は、段階的処理の制約を認めています。各タイルでは部分的な指数和を累積しますが、完全な正規化項は全タイル処理の完了まで利用できません。これは、従来のバッチ処理と比較したオンライン手法の一つの制約です。
しかし、「すべてのタイルを1回行った後、たとえば、すべてのn²タイルを行った後です。その時点で、softmaxを直接出力するのに必要なすべてのコンポーネントがあります」という説明は、この制約が実用的な問題にならないことを示しています。全タイル処理は単一パスで完了し、その時点で正規化に必要な全情報が利用可能になります。
「正規化項であるL3またはLのLのNを構築しました。私はこの最後のタイルの共有メモリに、そのようなものがすでにあります」という状況説明は、効率的な実装の詳細を示しています。最終タイルの処理完了時点で、累積された正規化項L_Nが高速な共有メモリ内に存在し、追加的なグローバルメモリアクセスなしに利用できます。
この設計の効率性は、再計算の回避にあります。従来の実装では正規化項計算のために全中間結果への再アクセスが必要でしたが、Flash Attention 2では段階的累積により、正規化項が処理の副産物として得られます。全タイル処理という必要な作業の完了と同時に、正規化に必要な全情報が準備されるのです。
27.4 V行列との最終乗算によるSoftMax出力
その時点で、これらのタイルのすべてを通した後、正規化項であるL3またはLのLのNを構築しました。私はこの最後のタイルの共有メモリに、そのようなものがすでにあります。
そして、それにより、指数化と除算を行い、それからすべてのコンポーネントを返すことができます。そして、最後に、タイル化されたV行列乗算でもう一度、そして、それは完全なsoftmaxを与えます。
このV行列との最終乗算によるSoftMax出力の説明は、Flash Attention 2の完全な処理フローの最終段階を示しています。正規化項の取得からSoftMax値の計算、そして最終出力の生成まで、一連の流れが明確に説明されています。
「指数化と除算を行い、それからすべてのコンポーネントを返すことができます」という処理は、標準的なSoftMax計算の最終段階です。正規化項L_Nが利用可能になった時点で、各要素について exp(xi - max) / L_N の計算を実行できます。重要なのは、この処理も高速な共有メモリ内で完結できることです。
この段階での効率性は、中間結果の適切な管理にあります。従来の実装では、SoftMax結果をグローバルメモリに書き戻してからV行列乗算のために再度読み込む必要がありました。しかし、Flash Attention 2では、SoftMax計算からV行列乗算まで、可能な限り共有メモリ内で連続処理します。
「最後に、タイル化されたV行列乗算でもう一次、そして、それは完全なsoftmaxを与えます」という最終段階は、再び確立されたタイリング技術の適用です。V行列も適切なタイルサイズに分割され、SoftMax結果との乗算が効率的に実行されます。この段階では、前述したKQ乗算と同様の最適化技術(メモリ合体、共有メモリ活用等)がすべて適用可能です。
完全な処理フローを振り返ると、Flash Attention 2は以下の統合されたパイプラインを実現しています:①KQタイル乗算→②オンラインSoftMax状態更新→③正規化項完成→④SoftMax値計算→⑤Vタイル乗算→⑥最終出力。各段階が効率的にタイル化され、グローバルメモリアクセスが最小限に抑えられています。
この統合設計により、従来の分離された処理段階(KQ乗算→SoftMax→V乗算)で発生していた無駄なメモリ往復が完全に排除されます。中間結果の大部分が高速メモリ内で処理され、最終的にO(n²)のメモリアクセスがO(n²/B)(Bはブロックサイズ)に削減されるのです。この最適化の集大成として、「完全なsoftmax」が従来とは比較にならない効率で実現されています。
28. 逆伝播での再計算戦略
28.1 n²サイズ活性化保存の回避必要性
逆伝播については取り上げません。タイル別にタイル別に再計算を行うことができ、それによりsoftmaxを保存することを避けることができます。softmaxを保存したい場合、それはすでにn²サイズのものですよね。n²の活性化を保存したくありません。逆伝播を行うときに、オンザフライでタイル別にタイル別に再計算する必要があります。
このn²サイズ活性化保存の回避必要性は、Flash Attentionにおける再計算戦略の根本的動機を明確に示しています。順伝播での最適化が逆伝播でも継続される必要性を具体的に説明しています。
「softmaxを保存したい場合、それはすでにn²サイズのものですよね」という指摘は、問題の規模を明確にしています。シーケンス長nに対して、注意機構のSoftMax出力はn×n行列となります。n=10,000の場合、100,000,000要素となり、FP16でも約200MB、FP32では約400MBのメモリが必要です。長いシーケンスでは、この数値は指数的に増大します。
この巨大なメモリ要求は、Flash Attentionの順伝播最適化の成果を台無しにしてしまいます。順伝播でO(n²)のメモリアクセスをO(n²/B)に削減したにも関わらず、逆伝播のために巨大な活性化を保存するのでは、全体的なメモリ効率改善が相殺されてしまいます。
「n²の活性化を保存したくありません」という方針は、一貫したメモリ効率戦略を示しています。現代GPUにおけるメモリボトルネックを考慮すると、大規模な中間結果の保存は性能に致命的な影響を与えます。特に、複数の注意層を持つ大規模Transformerでは、各層でn²サイズの活性化を保存することは実用的に不可能です。
この制約により、Flash Attentionは順伝播と逆伝播で一貫した設計哲学を採用する必要があります。順伝播でのオンラインSoftMaxと同様に、逆伝播でも中間結果の実体化を回避し、必要な値をオンザフライで再計算する戦略が不可欠になります。
この一貫性こそが、Flash Attentionを単なる順伝播最適化から、完全な訓練効率化システムへと発展させた重要な要因です。順伝播での革新が逆伝播でも活用されることで、訓練全体のメモリ効率とスループット向上を実現しているのです。
28.2 タイル別オンザフライ再計算実装
逆伝播を行うときに、オンザフライでタイル別にタイル別に再計算する必要があります。それ以外は、かなり標準的です。それは本当に勾配を計算するのと同じことです。ただし、タイル別にタイル別にその計算を行うだけです。
このタイル別オンザフライ再計算実装の説明は、Flash Attentionの逆伝播における技術的戦略を簡潔に示しています。順伝播で確立されたタイリング原理が、逆伝播でも一貫して適用されることを説明しています。
「オンザフライでタイル別にタイル別に再計算する必要があります」という要求は、逆伝播における根本的な実装方針を示しています。従来の逆伝播では、順伝播で保存された活性化(この場合はSoftMax結果)を直接利用していました。しかし、Flash Attentionでは、これらの値を保存する代わりに、必要な時点で即座に再計算します。
この再計算戦略は、前述した3重シグモイド関数の例で説明された再計算手法の直接的応用です。メモリアクセス削減のために計算コストを受け入れるという、現代GPU環境に最適化されたトレードオフを採用しています。200-300クロックサイクルのグローバルメモリアクセスを回避するために、数クロックサイクルの再計算を実行するのです。
「それ以外は、かなり標準的です。それは本当に勾配を計算するのと同じことです」という説明は、Flash Attentionの逆伝播が根本的に新しい数学理論を導入していないことを示しています。Chain rule、勾配の連鎖計算、誤差逆伝播の基本原理はすべて標準的なままです。革新は実装レベルでのメモリ効率化にあります。
「ただし、タイル別にタイル別にその計算を行うだけです」という実装方針は、順伝播との一貫性を示しています。各タイルが独立した処理ユニットとして機能し、必要な中間値(SoftMax結果、注意重み等)をタイル内で再計算し、勾配計算を実行します。グローバルメモリからの大規模データ読み込みは回避され、高速な共有メモリ内での処理が中心となります。
この設計により、逆伝播でも順伝播と同様のメモリ効率が実現されます。各タイルの処理完了時に、中間結果は即座に破棄され、次のタイルで必要に応じて再計算されます。この一貫したタイリング戦略により、Flash Attentionは訓練全体を通じて優れたメモリ効率を維持できるのです。
28.3 標準勾配計算のタイル化適応
それ以外は、かなり標準的です。それは本当に勾配を計算するのと同じことです。ただし、タイル別にタイル別にその計算を行うだけです。
この標準勾配計算のタイル化適応の説明は、Flash Attentionの逆伝播における重要な設計哲学を簡潔に表現しています。革新的な最適化でありながら、数学的基盤は確立された理論に基づいていることを明確にしています。
「それは本当に勾配を計算するのと同じことです」という表現は、Flash Attentionが新しい数学理論を発明していないことを示しています。Chain rule、偏微分、誤差逆伝播の基本原理はすべて変更されていません。∂L/∂Q、∂L/∂K、∂L/∂Vの計算は、数学的には従来の注意機構と完全に同一です。
この数学的同等性は重要な意味を持ちます。Flash Attentionで訓練されたモデルは、従来の実装で訓練されたモデルと理論的に同一の結果を生成します。近似や簡略化は一切行われておらず、純粋に実装効率の改善によって性能向上を実現しています。
「ただし、タイル別にタイル別にその計算を行うだけです」という実装方針は、最適化の核心を示しています。標準的な勾配計算を、タイル化された形式で実行することで、メモリアクセスパターンを根本的に改善しています。各タイルで必要な勾配成分を計算し、必要に応じて隣接タイルとの境界で適切な累積を行います。
この適応における技術的洗練は、勾配計算の依存関係をタイル境界と整合させることにあります。注意機構の勾配では、異なる位置間の相互作用が複雑ですが、適切なタイル分割とオンライン処理により、これらの依存関係を効率的に管理しています。
重要なのは、この標準勾配計算の保持により、Flash Attentionが既存の深層学習フレームワークと完全に互換性を持つことです。PyTorch、TensorFlowなどの自動微分システムは、Flash Attentionを通常の演算子として認識し、標準的な最適化アルゴリズム(Adam、SGD等)がそのまま適用可能です。
この設計選択により、Flash Attentionは革新的な性能向上を提供しながら、既存のエコシステムとの完全な統合を実現しています。研究者や開発者は、コードの大幅な変更なしに、劇的な性能改善を享受できるのです。
29. 講義の総括
29.1 ハードウェア理解が現代言語モデル成功の鍵
それで、今日の講義全体を総括します。ハードウェアは、今日私たちが持っているすべての言語モデルを本当に動かしているものの一種です。そして、本当にハードウェアを活用したい場合は、低レベルの詳細を理解する必要があります。
私が今日説明したすべての概念を、私は思うに、すべてのシステムの進歩が本当に関与していると思います。
このハードウェア理解の重要性に関する総括は、講義全体の核心的メッセージを明確に示しています。現代の言語モデルの成功が、単なるアルゴリズムの進歩ではなく、ハードウェアの深い理解に基づいていることを強調しています。
「ハードウェアは、今日私たちが持っているすべての言語モデルを本当に動かしているものの一種です」という表現は、現代AI革命の真の原動力を指摘しています。GPT、BERT、ChatGPTなどの画期的な言語モデルの実用化は、Transformer アーキテクチャの理論的美しさだけでなく、それらを効率的に実行するハードウェア最適化技術によって実現されました。
特に重要なのは、講義で説明されたGPU性能の超指数的成長(K20から H100まで)と、それを活用する最適化技術の発展です。テンソルコアによる行列乗算の劇的な高速化、Flash Attentionによる長いコンテキスト処理の実用化、メモリ階層の効率的活用など、これらすべてがハードウェア特性の深い理解に基づいています。
「本当にハードウェアを活用したい場合は、低レベルの詳細を理解する必要があります」という指摘は、表面的な知識では不十分であることを示しています。メモリアクセスパターン、バーストセクションアライメント、ワープ分岐、タイル化戦略など、これらの詳細こそが実際の性能差を生み出します。nanoGPTの事例(47次元追加で25%向上)は、この詳細への注意がいかに重要かを実証しています。
「私が今日説明したすべての概念を、私は思うに、すべてのシステムの進歩が本当に関与していると思います」という統括は、これらの技術が孤立した知識ではなく、現代AI システムの基盤となっていることを示しています。効率的なGPU利用、メモリ最適化、並列処理、数値精度管理など、これらすべてが統合されたシステムとして現代の言語モデルを支えているのです。
この理解は、将来のAI開発においても決定的に重要です。ハードウェアとソフトウェアの co-design、アルゴリズムのハードウェア適応、効率的な実装技術の習得など、これらの知識なしには次世代AIシステムの開発は不可能です。
29.2 システム進歩における低レベル詳細の重要性
私が今日説明したすべての概念を、私は思うに、すべてのシステムの進歩が本当に関与していると思います。私は思うに、すべてのシステムの進歩が、これらの概念の多くを本当に関与させます。
このシステム進歩における低レベル詳細の重要性に関する総括は、講義で説明された個々の技術が単独で機能するのではなく、統合されたシステムとして現代のAI進歩を支えていることを強調しています。
「すべてのシステムの進歩が本当に関与していると思います」という表現は、現代の高性能システム開発において、部分的な最適化では不十分であることを示しています。Flash Attentionの成功は、単一の技術的ブレークスルーではなく、タイリング、メモリ合体、再計算、低精度演算、オンラインアルゴリズムなど、複数の低レベル最適化技術の統合によって実現されました。
この統合的アプローチは、現代のシステム設計における新しいパラダイムを表しています。従来のソフトウェア開発では、抽象化により低レベル詳細を隠蔽することが重視されていました。しかし、性能が決定的に重要な現代のAIシステムでは、ハードウェア特性への深い理解と、それに基づいた緻密な最適化が不可欠になっています。
講義で説明された低レベル詳細の実例は、この重要性を具体的に示しています:
- メモリアクセスパターンの最適化:バーストセクションアライメント、メモリ合体
- データ構造の最適化:パディング、2の累乗アライメント
- 計算パターンの最適化:条件分岐回避、ワープ効率
- リソース管理の最適化:共有メモリ活用、レジスタ効率
これらの技術は、それぞれが個別に数%から数十%の性能向上をもたらしますが、適切に組み合わせることで、数倍から数十倍の劇的な改善を実現します。nanoGPTの47次元追加による25%向上、Flash Attentionによる長いコンテキスト処理の実用化など、これらはすべて低レベル詳細への注意の成果です。
「これらの概念の多くを本当に関与させます」という表現は、現代のシステム進歩が単一分野の専門知識では達成できないことを示しています。コンピューターアーキテクチャ、数値計算、アルゴリズム設計、メモリ管理など、複数の分野の深い知識を統合することで、真に革新的なシステムが生まれるのです。
この理解は、将来のAI研究者・開発者にとって決定的に重要です。高レベルなアルゴリズム設計だけでなく、ハードウェア特性を深く理解し、それに適合した実装技術を習得することが、次世代AIシステム開発の必須要件となっているのです。
29.3 GPU スケーリングとメモリ移動ボトルネックの関係
現在のGPUスケーリング、あなたが覚えておくべき種類のプロットです。それは本当にメモリ移動について考えることを奨励し、推進します。FLOPsをどのように削減するかだけでなく、本当にメモリ移動をより効率的にする方法について考える必要があります。
このGPUスケーリングとメモリ移動ボトルネックの関係についての総括は、講義の最も重要な洞察の一つを簡潔に表現しています。現代GPU最適化における根本的なパラダイムシフトを明確に示しています。
「現在のGPUスケーリング、あなたが覚えておくべき種類のプロットです」という言及は、講義で詳しく説明されたコンポーネント別スケーリング比較を指しています。このプロットは、計算性能(灰色線)の1〜100,000倍という驚異的成長に対して、メモリ帯域幅(緑線)とネットワーク接続性(青線)の相対的に緩やかな成長を示していました。
この不均衡こそが、現代GPU最適化の全ての戦略を決定しています。テンソルコアの導入により計算能力は爆発的に向上しましたが、DRAM技術の物理的制約により、メモリ帯域幅の改善は制限されています。結果として、多くの実際のワークロードで計算ユニットがデータ待機によりアイドル状態となり、高価なハードウェアが十分に活用されない状況が発生しています。
「それは本当にメモリ移動について考えることを奨励し、推進します」という表現は、最適化戦略の根本的転換を示しています。従来のCPU時代では、計算複雑度の削減が主要な目標でした。しかし、現代のGPU環境では、同じ計算内容でもメモリアクセスパターンを改善するだけで劇的な性能向上を実現できます。
この変化は、講義全体を通じて一貫して強調されたテーマです:
- Flash Attentionの成功:計算量O(n²)は変更せず、メモリアクセスをO(n²/B)に削減
- nanoGPTの事例:計算内容不変で、メモリアライメント調整により25%向上
- タイリング最適化:同じ行列乗算で、アクセスパターン改善によりt倍高速化
「FLOPsをどのように削減するかだけでなく、本当にメモリ移動をより効率的にする方法について考える必要があります」という指針は、新しい最適化思考の確立を求めています。演算数削減(FLOP reduction)から、データ移動効率化(memory movement optimization)への優先順位の転換です。
この理解は、将来の技術発展においても継続的に重要です。DRAMスケーリングの困難さを考慮すると、計算能力とメモリ帯域幅のギャップは今後さらに拡大すると予想されます。従って、メモリ効率を中心とした設計思想は、一時的なトレンドではなく、長期的な技術発展の方向性なのです。
この洞察を内面化することで、次世代の研究者・開発者は、ハードウェアの制約を深く理解し、それに適合した革新的なアルゴリズムとシステムを設計できるようになるのです。
29.4 Flash Attentionによる最適化技術統合の実証
そして最後に、特定の量の計算を行わなければならない場合、物事を最適化する方法は、データ移動を最適化することです。高帯域幅メモリまたはグローバルメモリからできるだけ多くの移動を避け、すべてを非常に非常に高速な共有メモリに配置し、それがタイリングやコアレッシング、再計算のような物事のflash attentionで良いパフォーマンスにつながります。
皆さん、ありがとうございました。
このFlash Attentionによる最適化技術統合の実証は、講義全体の学習内容が単なる理論的知識ではなく、実際の革新的システムにおいて体系的に活用される実例であることを示す完璧な結論です。
「特定の量の計算を行わなければならない場合、物事を最適化する方法は、データ移動を最適化することです」という原則は、現代GPU最適化の核心を表現しています。Flash Attentionの成功は、計算量の削減ではなく、同じ計算をより効率的なメモリアクセスパターンで実行することによって達成されました。O(n²)の本質的計算複雑度は変更せず、メモリアクセスをO(n²/B)に削減したのです。
「高帯域幅メモリまたはグローバルメモリからできるだけ多くの移動を避け、すべてを非常に非常に高速な共有メモリに配置し」という戦略は、Flash Attentionの設計哲学そのものです。200-300クロックサイクルのHBMアクセスを20クロックサイクルの共有メモリアクセスに置き換えることで、約10倍の性能向上を実現しました。
「それがタイリングやコアレッシング、再計算のような物事のflash attentionで良いパフォーマンスにつながります」という統合の説明は、講義全体の学習内容がFlash Attentionにおいてどのように結実しているかを示しています:
- タイリング:大きな注意行列を管理可能なブロックに分割し、各ブロックを共有メモリで効率的に処理
- コアレッシング(メモリ合体):連続メモリアクセスによりバーストモードを活用し、メモリ帯域幅を最大化
- 再計算:中間活性化の保存を回避し、必要時にオンザフライで値を再生成してメモリ使用量を削減
これらの技術は個別に適用されるのではなく、統合されたシステムとして機能します。オンラインSoftMaxアルゴリズムによりタイル内処理が可能になり、適切なメモリアクセスパターンによりハードウェア効率が最大化され、再計算戦略により逆伝播でもメモリ効率が維持されます。
Flash Attentionの実証的価値は、これらの「ランダムな切り離された事実」ではない最適化技術が、実際に現代AIシステムの基盤となっていることです。ChatGPTやGPT-4のような大規模言語モデルの実用化において、Flash Attentionのような効率化技術が決定的な役割を果たしています。
この統合実証により、講義で学習した詳細な技術知識が、単なる学術的興味ではなく、次世代AIシステム開発の実用的基盤であることが明確になります。理論と実装、ハードウェアとアルゴリズム、個別最適化と統合システム、これらすべての調和こそが、真に革新的な技術進歩を生み出すのです。