GPUの爆発的な性能向上

 (hazyresearch.stanford.edu)
1 ポイント 投稿者 GN⁺ 2024-05-13 | 1件のコメント | WhatsAppで共有
  • AI計算コストが増大する中、Hazy Researchは NVIDIA H100のtensor coreを休ませずに動かし続ける ことがGPU性能最適化の核心だと整理している
  • H100はhalf-precision行列積で 989 TFLOPs を出す一方、一般演算は約60 TFLOPsにとどまり、tensor coreが止まった瞬間に活用率が大きく落ちる
  • 最大性能に近づくには WGMMA、shared memoryの配置、アドレス生成、occupancyをあわせて扱う必要があり、wgmma.mma_asyncなしではマイクロベンチマークでピークの約63%にとどまる
  • 公開されたCUDA組み込みDSL ThunderKittens はtile・vector抽象化によってswizzlingやregister layoutのような複雑さを包み込み、FlashAttention系カーネルの作成を簡素化する
  • H100向けFlashAttention-2 forwardカーネルは約 100行 で書け、FlashAttention-2より約30%高速で、Based linear attentionカーネルは215 TFLOPsで動作する

H100性能を左右する条件

  • AIは多くの コンピュート を使い、Hazy Researchはここ数年、AIがより少ないコンピュートで動く、あるいは与えられたコンピュートでより効率的に動作するようにする取り組みを進めてきた
  • 実用的な目標は、GPUを高速化する中で得た知見を整理し、高速なカーネル作成を助けるCUDA組み込みDSL ThunderKittens を公開することにある
  • さらに広くは、ハードウェア理解が AIコンピュート の見方をどう変えたかを扱っている

NVIDIA H100の構造とボトルネック

  • H100 SXM GPUは次の構成を前提に議論されている
    • 80GB HBM3、帯域幅3TB/s
    • 50MB L2キャッシュ、帯域幅12TB/s、GPU全体で25MBセクション2つに分かれ、crossbarで接続されている

  • 132個のSM

    • 各SMは最大227KBのshared memoryを含む256KBのL1キャッシュを持ち、あわせて約33TB/sの帯域幅を持つ
    • Hopperの新ハードウェア Tensor Memory Accelerator(TMA) が非同期アドレス生成とメモリfetchを担う
    • 各SMは4つのquadrantで構成され、各quadrantにはwarp scheduler、512個のvector register、行列積用tensor core、並列組み込み命令群がある
    • すべてのコンピュートはSMで行われ、その大半は register 上で処理される
    • H100で性能を出す鍵はtensor coreを継続的にfed状態に保つことにある
    • H100はhalf-precision行列積で989 TFLOPs、「それ以外」の演算で約60 TFLOPsを提供する
    • tensor coreが使われるサイクルでは最低94%のハードウェア活用率に到達する
    • tensor coreが使われないサイクルでは最大でも6%の活用率にとどまる

WGMMA: 必須だが扱いにくい命令

  • H100には warp group matrix multiply accumulate 命令 wgmma.mma_async がある
    • PTXでは wgmma.mma_async
    • SASSでは HGMMA / IGMMA / QGMMA / BGMMA
  • 以前のGPUにおける wmma.mma.sync、mma.sync は、32個のthreadからなるwarp1つがtensor coreにデータを入れ、結果を待つ同期方式だった
  • wgmma.mma_async は連続した 128個のthread がSMの全quadrantにまたがって協調的に同期し、shared memoryから直接非同期の行列積を開始する
    • warpは行列積が進行している間、registerで別の作業を行える
    • 結果は任意のタイミングで待ち合わせできる
  • マイクロベンチマークでは、H100の総computeを引き出すためにこの命令群が必要だった
    • 使わない場合、GPUはピーク活用率の約 63% にとどまることが観察された
    • tensor coreがローカルリソース上でも深いハードウェアパイプラインを必要とするためかもしれない
  • 最大の難点は memory layout の複雑さにある
    • unswizzledなshared memory layoutはcoalescingが非常に悪く、L2帯域幅を多く要求する
    • swizzled layoutはドキュメントが誤っており、把握に時間がかかった
    • swizzled layoutは特定の行列shapeでしか動作しないように見え、wgmma.mma_async の他の機能とうまく噛み合わない
    • ハードウェアはtensor coreへ送る途中でsub-matrix transposeを行えるが、それはlayoutがswizzledでない場合に限られる
  • FlashAttentionのようなカーネルでは、TMAとL2キャッシュが十分高速なため、この問題をある程度隠せる
  • ハードウェアを完全に使い切るにはmemory requestをcoalescingし、bank conflict を避ける必要があるため、layout制御が重要になる

Shared memoryとbank conflict

  • Shared memoryのsingle-access latencyは約 30 cycles と見られ、この間にSMのtensor coreは32x32の正方行列積をほぼ2回実行できる
  • FlashAttentionのような従来の研究では主にHBM-SRAMボトルネックに注目しており、過去にはこのボトルネックが実際に重要だった
  • HBMが高速化し、tensor coreがチップの他部分より速いペースで伸びたことで、shared memoryの小さなlatencyさえも除去または隠蔽すべき対象になった
  • Shared memoryは32個のbankに分かれており、注意しないと bank conflict が起きる
    • 同じmemory bankに複数の異なるmemory片を同時要求すると、要求が直列化される
    • 経験上、これによってカーネルが不均衡に遅くなることがある
    • WGMMAやMMA命令が要求するregister layoutは、単純に書くとbank conflictを起こしやすい
  • 解決策は、複数の swizzling パターンでshared memoryを再配置し、conflictを回避することだ
  • 可能であればregisterとshared memoryの間の移動を避け、必要な場合はWGMMAやTMAのような組み込みハードウェアで非同期データ移動を行うほうが有利だ
  • 実際のwarpを使った同期移動は最も一般的ではあるが、ほぼ最悪のfallbackに近い

アドレス生成とTMA

  • H100はtensor coreもmemoryも高速なため、fetchする memory address を生成する作業そのものがチップ資源のかなりの部分を占める
    • 複雑なinterleaved patternやswizzling patternが加わると、さらに顕著になる
  • NVIDIAの Tensor Memory Accelerator(TMA) はglobal/shared memoryの多次元tensor layoutを指定し、そのtensorのsubtileを非同期にfetchしたうえで、完了時にbarrierをトリガーできるようにする
  • TMAはアドレス生成コストを下げ、pipeline構成も容易にする
  • TMAは wgmma.mma_async と同様に、H100の潜在力を引き出すうえで必須と評価されている
    • 経験上、WGMMAよりさらに重要かもしれない
    • register資源とinstruction dispatchを節約する
    • global memoryに対する非同期reduction機能もあり、複雑なbackwardカーネルで有用だ
  • TMAもまたswizzling modeを理解するには一部reverse engineeringが必要だったが、WGMMAほどつらくはなかった

Occupancyが隠してくれるコスト

  • CUDAで occupancy とは、同じ実行ハードウェア上でco-scheduledされるthread数を指す
  • SM quadrantのwarp schedulerは各cycleごとに、命令を受け取る準備ができたwarpへinstructionをissueしようとする
  • H100は前世代よりoccupancyへの依存が弱い面がある
    • 非同期機能のおかげで、単一のinstruction streamでもmemory fetch、matrix multiply、shared memory reduction、register mathを同時に忙しくできるためだ
  • それでもoccupancyは、ミスや同期コストを隠すうえで非常に有用だ
    • 完璧に設計されたpipelineは追加のoccupancyなしでも高速になりうる
    • 実際の観察では、NVIDIA GPUはoccupancyを念頭に置いて設計されているように見えた
    • 同期やミスの可能性が多いため、occupancyを高めると実効ハードウェア活用率が改善する場合が多かった
  • H100ではoccupancyは有用な水準だが、A100とRTX 4090ではそれぞれさらに重要になると見ている
    • H100に比べて同期instruction dispatchへの依存が強いためかもしれないと述べている

ThunderKittens: CUDA内の小さなDSL

  • ThunderKittens は、H100で高速なカーネルを簡単に書くために作られた CUDA組み込みDSL である
  • 当初は研究室内利用のために作られ、その後公開された
  • 名前は、kittensがかわいく、コードで kittens:: と打たせるのが面白いと思ったことに由来する
  • ThunderKittensは単純さを目標にしており、4つのtemplated typeを提供する
    • Register tiles: register file上の2D tensor
    • Register vectors: register file上の1D tensor
    • Shared tiles: shared memory内の2D tensor
    • Shared vectors: shared memory内の1D tensor
  • Tileはheight、width、layoutでparameterizedされる
  • Register vectorはlengthとlayoutでparameterizedされ、shared vectorはlengthのみを使う
    • shared vectorは一般にbank conflictを起こしにくい
  • 提供される演算はwarp levelまたは協調warp group levelでtile・vectorを操作する
    • initializer: shared vectorをzeroにする処理など
    • unary op: exp など
    • binary op: mul など
    • row/column op: row_sum など
  • ThunderKittensはCUDA内に組み込まれているため、Tritonのようなライブラリと違って抽象化が「gracefully」に失敗すると説明している
    • 足りない機能があれば、望む形で拡張できる

FlashAttentionの例と性能

  • ThunderKittensの例として、RTX 4090向けのシンプルなforward FlashAttentionカーネルが示されている
    • headdim=64のみを扱う
    • n は256の倍数である必要がある
    • 約 60行 のCUDAコードで書かれている
    • ハードウェア活用率は 75%
    • 複雑さの大部分はswizzling patternやregister layoutではなく、アルゴリズム自体にある
  • H100向けFlashAttention-2 forward passもThunderKittensで書かれている
    • TMA、WGMMA、swizzling mode、descriptorの複雑さをThunderKittensが包み込む
    • カーネルは約 100行
    • H100でFlashAttention-2より約 30% 高速
  • ThunderKittensは、GPU上で使える「mini-pytorch」のようにlayoutとinstructionを包み込み、primitiveを提供する
  • Based linear attentionや、今後公開される他アーキテクチャ向けのkernelもあわせて公開されている
    • Based linear attention kernelは 215 TFLOPs で動作する
    • アルゴリズム自体のrecomputeを考慮すると300 TFLOPsを超える
    • Linear attentionは理論上より効率的だが、実ハードウェアでは歴史的に効率が大きく低かった
    • この結果によって高スループットアプリケーションの適用範囲を広げられる可能性があると見ている

Tile中心の考え方

  • ThunderKittensがうまく機能した理由は、何でもしようとしなかったからだと見ている
    • CUDAはThunderKittensよりはるかに表現力が高い
    • ThunderKittensは小さく単純なDSLである
  • 核となる抽象化は small tile であり、これはAIとハードウェアが向かう方向に合っていると考えている
  • ThunderKittensは16未満の次元をサポートしない
    • ハードウェアもそのような小さい次元を特に求めていないと見ている
    • 「matrix multiplyが16x16より小さいなら、それをAIと呼べるのか確信できるか」という問題提起をしている
  • CPU時代の32-bit wordをregisterとして見る観点は、AIハードウェアには合わないと考えている
    • CUDAの1024-bit vector registerは正しい方向への一歩だと見ている
    • ここでのregisterは 16x16 tile のデータである
  • AIは依然としてmatrix multiply、reduction、reshapeが中心であるため、tile抽象化はAIにもハードウェアにも適していると考えている
  • 今後はAIのアイデアを、ハードウェアにうまくマッピングできる形へ再編成していく必要がある
    • recurrent stateのサイズはSMに収まる程度に大きくあるべきだ
    • compute densityはハードウェアが要求する水準を下回ってはならない
    • ハードウェアから学んだことをAI設計に合わせていくのが今後の重要な方向だ

AMD対応計画

  • ThunderKittensの AMD hardware 対応がまもなく提供予定

関連記事

1件のコメント

 
GN⁺ 2024-05-13
Hacker Newsのコメント

  • 「行列乗算が 16x16 より小さいなら、それは本当にAIなのか？」という問いが興味深い
    AIハードウェアの要件がだんだん明確になってきている。GPUはもともとまったく別の用途向けに設計されたが、行列乗算ハードウェアが優れていたためAIに使われ、「AI GPU」は実際のGPUにある一部の機能を削ぎ落とせる
    数値表現も16ビット浮動小数点、8ビット、2ビット、1ビットのように、より短くなっていく流れがあり、いずれ適切な落としどころが定まるはずだ。この記事は 16x16タイル を好むハードウェアがかなり理にかなっていることを示している。今この瞬間にも、誰かがすでにVHDLでこうしたものを書いているか、近いうちにそうする可能性が高い
    結局、「AI」の演算だけを、できるだけ不要なハードウェア上の重荷なしに実行する、より単純で、より汎用性が低く、安価なデバイスが出てきそうだ

    • GPUはすでに、可能な限り無駄のない AIマシン へと進化してきた。少なくともNervanaが設立された2014年から、GPUは古い技術なのでAIには不向きだと言われていたが、GPUがこれほど急速にAIマシンへ進化するとは予想されていなかったようだ
    • Appleはすでに数年前からこの方向に進んでいる。ダイ上の NPU はGPUやCPUとはまったく違う[1]
      Nvidiaもおそらく取り組んでいるだろうが、ゲーム／エンターテインメント／暗号資産／AIをひとまとめにしたデバイス、つまりビデオカードという形を維持するほうが、ビジネス上はよりよい選択なのかもしれない
      [1] https://github.com/hollance/neural-engine/blob/master/docs/a...
    • 「NVIDIAの嘘」というくだりが、競争の深さを物語っている。ドキュメントの誤りが完全な偶然であるはずはなく、図表は盗んだり写したりしやすいので、Nvidiaが意図的に残しておく効用があったのかもしれない
      Naveen RaoのNervanaが、Nvidia自身のドライバより速い Nvidia Maxwellドライバ を作っていた頃を思い出す。急成長する製品のドキュメント上のミスがすべて競争対策というわけではないが、研究者たちがwgmmaをリバースエンジニアリングするのに長い時間を要し、H100をめぐる米中の政治状況まで考えると、Nvidiaが堀を守るために以前の手口を使っているように見える
      だからH100の特異性を過度に掘り下げるよりも、「AIが求めるハードウェアとは何か」には商業的な事情も含まれると見るべきだ
    • AMDはすでに Versal ラインの第2世代まで来ている
      https://www.amd.com/en/products/accelerators/alveo/v80.html
      XDNA Architecture
      https://www.amd.com/en/technologies/xdna.html
    • Googleもこういうデバイスをもうほぼ10年作っているのでは？

  • 「NVIDIAの嘘。実際の128b swizzled wgmmaレイアウトについて、ひどく誤解を招く表現だ。この図のせいで取り戻せない人生の3週間を失ったので、公開の場でさらし者にする」というくだりが印象的
    AIの進展の非常に大きな部分が 行列乗算の最適化 のようなエンジニアリングにあり、そのエンジニアリングのかなりの部分が NVIDIAチップのリバースエンジニアリング だという事実に、驚く人がいるのか気になる

    • 構造そのものは大きな違いを生まない。十分に大きなデータで十分に大きなモデルを学習させれば、構造に関係なく似た結果になる傾向がある。だからAIの進歩の大半は、今では行列を非常に高速に掛けられるようになったおかげだと見なせる

  • ワープスケジューラ、4つの象限、Tensor Memory Accelerator、unswizzled wgmmaレイアウト…
    GPU用語と Star Trek風テクノバブル の境界がますます曖昧になっている

    • 記事を読みながらある程度は分かっていたが、「テンソルアクセラレータで象限をワープ中」と言うと、本当にStar Trekっぽい感じがする
      他の記事を見ているときも、たまにこう思うことがあった。誰かがここの記事リンクを受け取って読んだら、どんな感じだろうかと。ワープコアについて議論するTrekオタクのイベント会場に入り込んだような感じになりそうだ
    • この言葉を見て一歩引き、用語を新しい目で見ることになり、本当にその通りだと笑ってしまった

  • AI推論の消費電力を減らし、速度を上げるには、アナログ近似回路へ移行するのが最善のように思う
    完全な浮動小数点の乗算と加算が必要なのではなく、2つの入力電圧を受け取って、乗算結果に十分近い出力電圧を出す装置が必要なだけ

    • この方向で取り組んでいる人を知っているが、大きな難関は、既存のチップ製造技術でアナログ論理が可能なものを作る方法、アンテナのように動作しない設計、そして物理チップごとに製造公差が異なるため、実行するモデルをチップごとに微調整しなければならない可能性だと聞いた
      大きな利点は、float16を16本の線で表現する代わりに、1本の線の電圧でその数を表現すること。理論上はfloat32よりはるかに高い精度も可能かもしれない。また、値を算術論理ユニットにロードせず直接接続できるので、ダイ面積と消費電力の削減が潜在的に桁違いの規模になり得る
    • アナログ回路が実用的に有用になるにはまだ遠いと思うが、不正確さを許容できる場所ではノイズのあるデジタル回路かもしれない
      例えば、出力ビット100万個のうち1つが反転することを受け入れ、性能/電力比を改善するような形。単一の無限大値が全体を壊しかねないfloat32では難しいだろうが、int8では0が欲しかったのに時々128が出る程度なら耐えられそうに思う
      [1] H100の行列浮動小数点ユニットが実際にIEEE 754に準拠しているのかはよく分からない
    • さらに一歩進めると、実際の生物学的な脳が動作する仕組みに似ていながら、容易に生産できる何かが必要だと思う
      生物学的ニューラルネットワークは、一般的な人工ニューラルネットワークのように完全結合に近いものではなく、ニューロンの入出力接続係数は10未満で、非常に局所的。生物学には、私たちの知る限りバックプロパゲーションもなく、代わりにフィードバックと再帰がある
      まだ分かっていない中枢神経系の機能に不可欠な補助細胞やプロセスがあるかもしれない。高レベルでも相当量の「ハードコードされた」結合性がある可能性があり、すでに一部は知られている。例えば耳の聴覚ニューロンは接続されており、音の位置を特定するために畳み込みに似たことが起きる。これは創発現象ではなく、訓練なしでも可能な機能
      生命は数十億年と同程度の世代数を経てこれを見つけ出したのだから、驚くことではない。理論的にはソフトウェアでも可能だろうが、霊長類/人間の脳にある1兆個以上のニューロンを考えると、今日の千コア級マシンでも極めて難しい。「クラウド」だとしても、必要な結合性とレイテンシを満たせないだろう
      こうしたアプローチでワームや昆虫くらいをうまくモデル化できたら面白そう
    • 十分な範囲と精度を同時に満たすのは、ほぼ不可能に見える
    • 正直、デバッグするには悪夢のように見える

  • この記事を見て、CS 149の並列プログラミングの授業で感じた楽しさを思い出した

    • KayvonとKunleはすごい。2学期前にCS149 Parallel Programmingを受講したけど、本当に良かった :)

  • この記事の文体は本当に印象的で、AMD MI300xでこれを見るのが楽しみ。私の機材で時間を使ってみたいなら知らせてほしい

    • AMD製品でAI作業をたくさん試したことがあるのか気になる。RTX 4090に2500ドル以上は使いたくないので、実験用や入門用としてRX 7900XTXを検討している
      実際どれくらいうまく動くのか、あるいはもう少し貯めて7900 XTではなくXTXを買うほうがいいのか、VRAMが減ると実用性にどれほど影響するのかも気になる
    • 良い文章は明確で曖昧でないべき。話し言葉なら途中で遮ってclarificationを求められるが、文章にはメッセージを伝える機会が一度しかない
      読者が著者たちの言おうとしていることを知るためにknowyourmeme.comまで調べに行かなければならないようではいけない。このタイトルが何を意味しているのかさえ分からず、それだけ目標を大きく外していると思う
    • 本当に？ Wallstreetbets時代のPTSDを思い出す

  • こういう記事を完全に理解するには、どこから始めればいいのか、どんなロードマップをたどればいいのか気になる

    • GPUプログラミングを学ぶのに良い講義がある。4.0講義あたりで必要な基礎が得られる: https://youtube.com/playlist?list=PLzn6LN6WhlN06hIOA_ge6Srgd...
      それから、ベクトル-行列乗算を行うCUDAカーネルを自分で書いてみるとよい。pycudaを使えばカーネルに集中し、残りはPythonで書ける。ChatGPTに、4000要素のベクトルと4000x12000行列を掛ける実装を自分で作りたいと伝え、全工程を案内してもらえばよい
      GPUレンタルはRunpodがよく、今は低価格GPUからH100まである。最初は低いグレードのGPUから始めればよい
    • 深く入るなら、Spiralの行列乗算プレイリストを見るとよさそう: https://www.youtube.com/playlist?list=PL04PGV4cTuIWT_NXvvZsn...
      Spiralで行列乗算カーネルを実装して最適化するのに2か月を費やした

  • GitHub READMEのグラフ(https://github.com/HazyResearch/ThunderKittens/blob/main/att...)があまりにも目が回る。こういう波打つ棒グラフって合法なのか？ :P

    • 同意。何か錯視を仕込もうとしているように見える。むしろ棒なしで数字だけ見たほうがよさそう
    • matplotlibのxkcdテーマのように見える[1]。それでも波打ちがやりすぎだという点には同意
      [1]: https://matplotlib.org/stable/gallery/showcase/xkcd.html#sph...

  • ThunderKittensという名前が素晴らしい。ThunderKittensが、順伝播より一桁難しい FlashAttentionの逆伝播 を扱うところを見てみたい

  • こうした研究は、今日すでにNPUを作っているチームがやってきたことではないのか？ たとえばGroqチップはAI専用の構造を使っているから、今のような性能を出せている。コンシューマー向けでは Apple Silicon もかなり有能だ
    この分野の人間ではないが、比較的遅い経路で通信する汎用プロセッサだけでは限界があるように思う。ハードウェアレベルで設計を考え直し、最終的にはコンシューマー市場で価格を下げる方向のほうが、長期戦略としてはより良さそうに見える

    • コンシューマー向けでApple Siliconがかなり有能だというのは、あまり確信が持てない。Redditのlocalllamaサブレディットを見ると、実用的な速度を出そうとして苦労している CPUユーザー の不満投稿が多い
      数百ドルでNvidia GPUを買うことも、900ドルで4050・6GB VRAM搭載のゲーミングノートPCを買うこともできるのに、CPUベースのAIを有能と呼ぶのは難しい
      職場にもGPUがなく、CPUベースで試してみたが、小さなモデルを使って待つ以外に現実的な選択肢はなかった。結局GPU搭載マシンを依頼することになった
      「技術的には可能」と「実際に快適に使える」は別物だ。Nvidiaは本当に使いやすく、CPUは苦痛でストレスがたまった