GPUの爆発的な性能向上

 (hazyresearch.stanford.edu)
1 ポイント 投稿者 GN⁺ 2024-05-13 | 1件のコメント | WhatsAppで共有

H100 GPUの特徴

  • 80GB HBM3メモリと3TB/sの帯域幅を提供（実際にはやや低い）
  • 50MBのL2キャッシュと12TB/sの帯域幅を提供。GPU上で2つの25MBセクションに分割され、クロスバーで接続されている（クロスバーは性能低下要因）
  • 132個のStreaming Multiprocessor（SM）で構成され、各SMは以下の構成を持つ:
    • 256KB L1キャッシュ内に最大227KBの共有メモリを提供（合計で約33TB/sの帯域幅）
    • 非同期アドレス生成およびメモリフェッチが可能なTensor Memory Accelerator（TMA）を提供。オンチップメモリネットワーク対応などの機能もあるが、この投稿では扱わない
    • 4つのクアドラントに分かれ、各クアドラントはwarp scheduler、512個のvector register（それぞれ32個の4-byteワードを含む）、行列積用のtensor core、sum/multiplyなどの並列演算を支える内蔵instructionで構成される

H100 GPU性能最適化のヒント

  • Tensor Coreを最大限活用することが重要。H100は989 TFLOPsのFP16行列積演算性能を持つため、Tensor Coreの活用度によってGPU全体の利用率が大きく左右される
  • ただし、Tensor Coreを最大限活用するには以下を考慮する必要がある:
    • Warp Group Matrix Multiply Accumulate（WGMMA）命令の活用が必須だが、扱いが難しい
    • Shared Memoryは思ったほど速くなく、利用時には細かな注意が必要
    • アドレス生成のコストが高いため、Tensor Memory Accelerator（TMA）などを使って最適化する必要がある
    • Occupancyを高めることは依然として有効であり、registerが主要リソースとなる
  • これらの特徴はH100だけでなく他のGPUにもある程度当てはまるが、特にH100ではTensor Coreの活用が重要で難しい部類に入る

ThunderKittens : H100に最適化されたCUDA埋め込みDSL

  • NVIDIA H100など最新GPUの性能を最大限引き出すために開発された、CUDAベースの埋め込みDSL
  • 以下の4つのタイルベースのテンプレート型を提供:
    • Register Tile（レジスタに保存された2Dテンソル）
    • Register Vector（レジスタに保存された1Dテンソル）
    • Shared Tile（Shared Memoryに保存された2Dテンソル）
    • Shared Vector（Shared Memoryに保存された1Dテンソル）
  • さらに、タイル操作のためのさまざまな演算子（exp、mul、sumなど）をwarpまたはwarp groupレベルで提供
  • 既存のFlash Attention、Flash Attention 2カーネルをThunderKittensで実装すると、コードが大幅に簡潔になり、H100で最大30%の性能向上
  • Based Linear AttentionカーネルもThunderKittensで実装し、215 TFLOPsの性能を達成（アルゴリズムの特性上、recomputeを含めると300 TFLOPs以上）

哲学的観点からの考察

  • ThunderKittensがうまく機能する理由は、あらゆるものをサポートしようとせず、シンプルでGPUアーキテクチャにうまく適合するタイルベースの抽象化を提供しているため
  • 従来の32-bitワードの代わりに1024-bit vector registerを使うのも進歩だが、16x16タイルをレジスタの単位として捉えるパラダイム転換が必要
  • AIワークロードは結局のところ行列積、reduction、reshape中心であることを踏まえると、タイルベースのアプローチは妥当であり、ハードウェアの観点でもsystolic array以外に小規模な行列積を支援する方向へ進化していくだろう
  • さらに進んで、ハードウェアに最適化された形でAIアルゴリズムを設計する方向へ思考を転換する必要がある。たとえばRNNの状態サイズをSMに収まる程度に制限し、演算密度もハードウェアが要求する水準に合わせるといったアプローチが必要

GN⁺の意見

  • ThunderKittensはCUDAに慣れた開発者にとって魅力的な選択肢になり得る。既存のカーネルコードを大きく変更せずに、容易に性能向上を得られるため
  • ただし初心者にとっては依然として参入障壁が高い可能性がある。今後はより多様なサンプルコードや学習資料による支援が必要そうだ
  • H100だけでなくAMDなど他のGPUにもThunderKittensの対応を広げる計画は興味深い。ベンダーロックインの低減に寄与する可能性がある
  • 最終的にはAIモデルやアルゴリズムそのものをハードウェアに最適化された形で設計することが非常に重要なポイントとなる。そのためにはハードウェア特性への深い理解が前提となるが、ThunderKittensは開発者がそうしたインサイトを得る助けになり得る
  • Hazy Researchの継続的な研究開発とオープンソースへの貢献は、CUDAエコシステムの活性化に大きく役立つと期待される。ただし長期的には、より抽象化レベルの高いフレームワークの登場も必要に思われる

関連記事

1件のコメント

 
GN⁺ 2024-05-13
Hacker Newsの意見

  • AIハードウェアの要件はますます明確になってきている。GPUはもともと別の目的のために設計されたが、優れた行列積ハードウェアを備えているためAIに使われている。"AI GPU"では実際のGPUの一部機能を省ける可能性があり、より短い数値（16ビット、8ビット、2ビット、1ビット浮動小数点）へ向かう傾向がある。この論文は、16x16タイルを好むハードウェアに多くの利点があることを示唆している。

  • AI専用の補助プロセッサ（NPU）が必要である。特に開発者、専門家、ゲーマーなど向けのプロシューマー級デスクトップシステムで必要だ。GPUは企業では機能するが、パーソナルコンピューティングの観点ではAI用途に使うには扱いにくい。特にVRAMの制限と、Vulkan以外に標準的なオープンAPIが存在しないことが問題だ。

  • AI研究を前進させるには、神経科学や心理学をより深く研究する必要がある。また、ニューラルネットワークのグラフトポロジーに関することも関係している可能性がある。

  • CUTLASSをユーザーフレンドリーにしたものなのか？

  • ThunderKittensのマスコットには、猫／ソニーAIBOのような雰囲気がある。AIでうまく生成されたように見える。

  • ユニバーサル・ベーシック・コンピュート（UBC）がユニバーサル・ベーシック・インカムの代替として検討されるなら、それは非常にディストピア的な未来になるだろう。NVIDIAのような一社があらゆる計算資源を作ると想像してみてほしい。