世界最速のスパコンFrontier、搭載GPUのわずか8%でGPT-4並みの1兆パラメータモデルを実行

オークリッジ国立研究所に設置されているFrontierスーパーコンピューターを使って、研究者らはChatGPT-4に匹敵する1兆パラメータの大規模言語モデル（LLM)を学習させた。同チームは新たな論文でその方法を説明しているが、同時にその過程でいくつかの課題にも直面したようでこれについても語っている。

Frontierスーパーコンピューターは、世界をリードするスーパーコンピューターであり、2024年1月現在稼働している唯一のエクサスケール・マシンだ。このスーパーコンピューターは、AMDの第3世代EPYC Trento CPUとInstinct MI250X GPUアクセラレーターを搭載して一から設計されている。米国テネシー州のオークリッジ国立研究所（ORNL）に設置され、エネルギー省（DOE）が運用している。現在、8,699,904コアを使用して1.194エクサフロップ/秒を達成している。HPE Cray EXアーキテクチャは、HPCとAI向けに最適化された第3世代AMD EPYC CPUと、AMD Instinct 250Xアクセラレータ、Slingshot-11インターコネクトを組み合わせている。FrontierはTop500.orgのスーパーコンピューター・リストで1位を維持し、その優位性を示している。

Frontierが達成した新記録は、LLMを訓練し、搭載ハードウェアを最も効率的に使用するための効果的な戦略を実行した結果である。同チームは、220億、1750億、1兆のパラメータを広範にテストした結果、注目すべき結果を達成することができた。しかし、チームは1兆個のパラメーターを持つLLMを訓練するのに3,072個のGPUしか使わず、1,750億個のパラメーターを持つ別のLLMを訓練するのに1,024個しか使わなかったと報告している。MI250Xは、比較的旧式のハードウェアであるにもかかわらず、その実力を発揮したのだ。

さらに興味深いのは、Frontierのスーパーコンピューター全体が37,000個のMI250X GPUを搭載しており、GPUプール全体を使ってLLMを動かした場合のパフォーマンスが想像できることだ。AMDはまた、AI性能をさらに加速させる堅牢なROCm 6.0エコシステムを備えた、まったく新しいスーパーコンピューターにMI300 GPUアクセラレーターを実装しようとしている。

220億、1,750億、1兆のパラメーターに対して、それぞれ38.38%、36.14%、31.96%のGPUスループットを達成しました。1,750億パラメータモデルと1兆パラメータモデルの学習では、それぞれ1024個と3072個のMI250X GPUで100%の弱いスケーリング効率を達成しました。また、これら2つのモデルでは、89%と87%の強力なスケーリング効率を達成しました。
arXiv

この論文では、このような大規模なLLMを訓練する際の重要な課題は、必要なメモリ量であり、最小でも14テラバイトであったと指摘している。このため、それぞれ64GBのVRAMを搭載した複数のMI250X GPUを使用する必要があったが、これには並列性という新たな問題が発生した。より多くのGPUをLLMに投入するには、より多くのリソースを実際に効果的に使用するために、ますます優れた通信が必要になる。そうでなければ、余分なGPU馬力のほとんど、あるいはすべてが無駄になってしまう。

彼らはMegatron-DeepSpeedやFSDPのようなフレームワークを反復し、トレーニング・プログラムがフロンティア上でより最適に実行されるように変更したということだ。つまり、ワークロードのサイズが大きくなるにつれて、より多くのGPUが可能な限り効率的に使用されるようになったということだ。

一方、強スケーリング効率は、1,750億パラメータLLMで89%、1兆パラメータLLMで87%と、わずかに低かった。強スケーリングとは、ワークロードのサイズを変えずにプロセッサ数を増やすことを指し、アムダールの法則によれば、コア数が多くなると有用性が低くなる傾向がある。

しかし、研究チームは、Frontierでこの効率を達成するためには、いくつかの問題があることを指摘し、「AMD GPUでの効率的なトレーニング性能を探求する研究がもっと必要であり、ROCmプラットフォームはまばらである」と述べている。論文にあるように、この規模での機械学習のほとんどは、NVIDIAのCUDAハードウェア・ソフトウェア・エコシステムの中で行われており、AMDやIntelのソリューションは比較にならないほど未発達であり、両社はNVIDIAのソフトウェアソリューションに追いつく必要がある。当然ながら、このような取り組みは、これらのエコシステムの発展を促進するだろう。

Frontierが現在採用しているハードウェアは、業界では比較的新しいものではない。生成AI分野での継続的な進歩に伴い、市場が今後より多くのコンピューティング・パワーを必要とすることは明らかであり、だからこそ、この分野向けに設計されたハードウェアの進歩が次世代の進歩に不可欠なのである。

論文

arXiv: Optimizing Distributed Training on Frontier for Large Language Models

参考文献

Wccftech: AMD-Powered Frontier Supercomputer Uses 3K of Its 37K MI250X GPUs To Achieves a Whopping 1 Trilllion Parameter LLM Run, Comparable To ChatGPT-4

研究の要旨

大規模言語モデル（LLM）は、基礎モデルとして顕著な成功を収めており、微調整を通じて様々な下流アプリケーションに恩恵をもたらしている。最近のロス・スケーリングに関する研究では、より大きなLLMが、より小さなLLMに比べて優れた性能を発揮することが実証されている。とはいえ、何十億ものパラメータを持つLLMをトレーニングすることは大きな課題であり、かなりの計算資源を必要とする。例えば、20兆個のトークンに対して1兆個のパラメータを持つGPTスタイルのモデルをトレーニングするには、1億2,000万エクサフロップスという驚異的な計算量が必要だ。この研究では、オープンサイエンスに特化した世界初のエクサスケール・スーパーコンピューターであるFrontierからこの計算を引き出すための効率的な分散トレーニング戦略を探求します。テンソル並列、パイプライン並列、シャーデッドデータ並列など、様々なモデル・データ並列トレーニング技術を駆使し、Frontier上で1兆パラメータモデルのトレーニングを容易にする。これらの手法と関連するパラメータを実証的に評価し、メモリフットプリント、通信レイテンシ、GPUの計算効率に与える影響を明らかにする。これらの手法間の複雑な相互作用を分析し、ハイパーパラメータのチューニングを通じて高いスループットを達成するために、これらの手法を組み合わせる戦略を見出した。経験的分析とハイパーパラメータチューニングにより、様々なサイズの大規模LLMをトレーニングするための効率的な戦略を特定した。220億、1750億、1兆のパラメータに対して、それぞれ38.38%、36.14%、31.96%のGPUスループットを達成しました。1,750億パラメータモデルと1兆パラメータモデルのトレーニングでは、それぞれ1024個と3072個のMI250X GPUで100%の弱いスケーリング効率を達成した。また、これら2つのモデルでは、89%と87%の強力なスケーリング効率を達成した。