近年、大規模言語モデル(LLM)の台頭により、機械学習タスクにおけるGPU活用の効率化が求められています。しかし、GPU性能を正確に評価することは大きな課題となっています。

従来のGPU使用率の限界

一般的に用いられるGPU使用率指標には問題があります。nvidia-smiなどのツールで確認できるこの指標は、実際の計算効率を正確に反映しないのです。驚くべきことに、計算を行わずにメモリの読み書きだけでも100%のGPU使用率を示すことがあります。

この事実により、機械学習分野でのパフォーマンス指標や評価手法の見直しが進んでいます。研究者たちは、LLMの学習や推論タスクにおけるGPU性能をより正確に測定し最適化する方法を模索しています。

新たな指標:モデルFLOPS使用率(MFU)

GPU使用率の限界に対処するため、研究者たちは代替指標を導入しました。その一つが、GoogleのPaLM論文で紹介されたモデルFLOPS使用率(MFU)です。MFUは観測されたスループットと理論上の最大スループットの比率を測定し、GPU性能をより正確に表現します。

しかし、MFUにも欠点があります。計算が複雑で、パラメータやフレームワークに依存するのです。それでも、MFUはGPU使用率と実際の計算効率の間に大きな乖離があることを明らかにしました。例えば、GPU使用率100%のLLM学習でも、MFUは20%程度にとどまることがあります。

これは、一般的なLLM学習の35-45%という範囲を大きく下回っています。この発見は、GPU性能指標をより深く理解する必要性を浮き彫りにしました。

Trainy AIの取り組み

GPU クラスター管理インフラを専門とする Trainy AI の研究者たちは、ある基盤モデル企業の LLM 学習効率を最適化する課題に取り組みました。彼らのアプローチは、PyTorch 向けに一般的に推奨される一連のパフォーマンスチューニング技術を実装することでした。

最適化の内容

これらの最適化には以下が含まれます:

  1. データローダーのパラメータ調整によるGPUの飽和
  2. 混合精度学習によるテンソルコア使用率の最大化
  3. apexやdeepspeedからのfused optimizerの採用
  4. 学習タスク向けに設計されたインスタンスとネットワーキングの利用

これらの手法を適用することで、Trainyは100%のGPU使用率と大きな電力消費を達成し、一見パフォーマンスが向上したように見えました。

しかし、実際の計算効率をより包括的に理解するため、チームはさらに一歩踏み込んで学習ワークロードのMFUを計算しました。これは、GPU使用率のみに頼ることの限界を認識したためです。

GPU アーキテクチャの理解

GPU使用率が性能指標として限界がある理由を理解するには、GPUアーキテクチャの知識が重要です。GPUはコアとマルチプロセッシングマネージャ(NVIDIAではSM、AMDではCU)で構成されています。

例えば、GH100 GPUには144個のSMがあり、各SMが複数のCUDAコアを管理しています。NVIDIAによるGPU使用率の定義は曖昧ですが、DatadogのNVMLドキュメントはより明確な説明を提供しています。

しかし、この指標は誤解を招く可能性があります。GPUの活動を示すだけで、計算効率を反映していないのです。CUDAカーネルが起動すると、SMによってコア全体に作業が分散されますが、使用率の割合はこれらの計算の強度や効果を反映していません。

PyTorch Profilerによる分析

パフォーマンスのボトルネックをさらに調査するため、研究者たちはPyTorch Profilerを使用してモデルの学習ループをプロファイリングしました。この分析により重要な洞察が得られました。

Softmaxカーネルが高いGPU使用率を示す一方で、SM(ストリーミングマルチプロセッサ)効率が低いことが判明したのです。この不一致は懸念を引き起こしました。単純なSoftmax実装が大規模言語モデルのボトルネックとなることはよく知られているからです。

低いSM効率は、高いGPU使用率にもかかわらず、モデル実行に非効率な部分がある可能性を示唆しています。この観察結果は、GPU使用率のみに頼ることの限界を裏付けています。

このようなメモリバウンドな操作に対処するため、FlashAttentionなどの様々なカーネルフュージョン技術が開発されてきました。プロファイリング結果は、LLM学習の最適化にはより繊細なアプローチが必要であることを強調しています。GPU使用率だけでなく、SM効率の向上に焦点を当てる必要があるのです。

SM効率の重要性

SM効率(SMアクティビティとも呼ばれる)は、NVIDIA GPUの重要な指標です。これは、一定の時間間隔内でアクティブなSMの割合を測定します。例えば、NVIDIA H100 GPUには132個のSMがあり、各SMが128個のコアを管理し、合計16,896個のコアがあります。

この指標は、CUDAカーネルが利用可能なSMをどれだけ効果的に使用しているかを示します。例えば、H100上で10秒間連続して実行されるCUDAカーネルが1つのSMしか使用しない場合、100%のGPU使用率を示しますが、SM効率はわずか0.7%です。

この不一致は、GPU使用率だけでなく、他の指標も考慮することの重要性を浮き彫りにしています。SM効率を層ごとに監視することで、研究者はLLM学習における最適化の機会を特定し、より的確なパフォーマンス改善と計算効率の正確な評価が可能になります。

LLM学習の最適化手法

研究者たちは、トランスフォーマーブロック内の層を融合することに焦点を当てました。この手法では、PyTorch のネイティブな層定義を CUDA や Triton で実装された GPU カーネルに置き換え、複数の層を単一のカーネルに統合します。

最適化の対象には以下が含まれます:

  1. Softmax (Flash Attentionを使用)
  2. MLP
  3. ドロップアウト層の正規化残差加算操作

これらの融合カーネルは、Flash Attentionなどのライブラリで提供されており、パフォーマンスの向上とメモリ使用量の削減を実現します。

実装の課題

実装における主な課題は、置き換えに適した層を特定することでした。torch.compileの自動最適化が、FSDPなどの新しい分散戦略と互換性がなかったためです。これらの制限により、融合カーネルの手動実装が必要となりました。

最適化の成果

最適化の取り組みは大きな成果を上げました:

  1. 学習時間が4倍高速化
  2. モデルFLOPS使用率(MFU)が20%から38%に向上

これらの改善は、融合カーネルの実装と、利用可能な3.2 Tbps Infinibandインフラを効果的に活用するためのモデル並列性の微調整によるものです。

パフォーマンス測定の推奨事項

この研究では、GPUクラスターのパフォーマンスを正確に測定するため、SM効率とGPU使用率の両方を追跡することを推奨しています。GPU使用率はマシンがアイドル状態かどうかを示し、SM効率はGPUがどれだけ効果的に使用されているかを示します。

MFUの計算は有益ですが、継続的なモニタリングには複雑すぎる場合があります。NvidiaのData Center GPU Manager (DCGM)はデフォルトでSMアクティビティを追跡します。SM占有率などの他の指標は、各SMのワークロードに関する詳細な洞察を提供しますが、解釈がより複雑です。

より深い理解を得るには、Pytorch Profilerブログ、DCGMドキュメント、Nsightのプロファイリングガイドを参照することをお勧めします。