GPUを使ったCerebras-GPTの学習と推論

こんにちは。CerebrasプリセールスエンジニアのNakadaです。
今回はCerebras-GPTのみならず、LLM(大規模自然言語モデル)を利用するにあたって、どのくらいの計算リソースが必要になるかをCerebras-GPTを使って確認しましたので結果を共有いたします。

これまでのCerebras-GPTについてのブログは以下に掲載していますので、ご興味がある方はご覧ください。

前々回のブログ：Cerebras-GPTがリリースされました

前回のブログ：Cerebras-GPTでファインチューニング

LLMには、Transformer技術をベースとした様々なモデルがあります。共通して言えることはTransformerを利用したエンコード層もしくはデコード層を積み重ねる事で精度が上がるということです。1750億パラメータを持つGPT3(ChatGPTのベースモデル)では96層のTransformerデコード層を実装しているそうです。

そのため、LLMを利用する際には学習・推論ともに巨大なパラメータをいかに高速に処理するかを考えなければなりません。

今回は、Cerebras-GPTモデルをDGXA100で学習及び推論させ、どのくらいの時間が必要だったかを共有し、みなさんが今後LLMを利用する際にどのくらいのGPUリソースが必要になるかイメージすることが出来れば光栄です。

Cerebras-GPTの学習時間

Cerebras-GPTモデルは現時点でパラメータサイズの異なる7つのモデルがリリースされています。今回はその中から、2.7B、6.7B、13Bの3つモデルを利用し、データセットとしては、一般に公開されているWikipedia日本語コーパスでファインチューニングを行いました。

※Cerebras-GPTのモデル名に使われている「B」はBillion(10億)の頭文字で、モデルのパラメータサイズを表します。2.7Bは27億パラメータ、6.7Bは67億パラメータ、13Bは130億パラメータを持ったモデルとなります。

まず初めに、これら3つのモデルは650億パラメータを持つLLaMAや1750億パラメータを持つOpenAI GPT3などのLLMと比べると小さいモデルとなりますが、画像処理で使われるUNETやRESNETなどのモデルは以下のグラフで比較したように10億パラメータ以下のモデルが多く、それらの画像系モデルと比べるとCerebras-GPTモデルは圧倒的に大きなモデルとなります。

※画像系モデルとLLMのパラメータ数

この規模のモデルを実行する場合、前回のファインチューニングの際に説明しましたが、2.7B以上のモデルはNVIDIAのA100(40GB)ではメモリに乗り切らず、Out Of Memoryが発生し学習することもできませんでした。そのため、今回はMicrosoft社が提供しているDeepSpeedのライブラリの中からZeRO、及びZeRO-Offloadを使うことにします。ZeROを用いることで複数のGPU間でメモリを効率的に使用できるようになり、ZeRO-Offloadを用いることで学習パラメータ等のデータをCPUメモリにオフロードすることができます。これらを併せることでOut Of Memoryを回避して学習することが可能になります。DeepSpeedについては別の機会にブログにて掲載しようと思います。

【学習設定情報】

結果は以下となりました。8枚のGPUでデータパラレルをしても2.7Bで約2週間、13Bでは約2か月半必要でした。なお、今回DeepSpeedを利用してGPUのメモリ不足は解消しましたが、バッチサイズを大きく設定することはできず、最小の1GPU 1バッチという設定となりました。こちらについてはDeepSpeedの更なるチューニングやA100 80GBモデルを利用するなどで、改善は可能です。

【学習に必要な時間】

上記の結果ですが、みなさんの感想は如何でしょうか。おそらく、LLMを利用する際にGPUの確保をどのようにするか頭を悩ませているのではないでしょうか。なお、今回は1GPUにCerebras-GPTモデルを展開し、8つのGPUでデータパラレルを行っています。更にGPUを増やすことでパフォーマンスを向上させることが可能です。

Cerebras-GPTの推論時間

続いて、推論時間の測定です。Cerebras-GPTの推論は文書生成となります。プロンプトを与え、生成が完了するまでの時間を測定しました。なお、Cerebras-GPTの7モデルで実行しました。

推論時には、学習済みモデルを利用しますが、もちろん利用するモデルサイズによって利用されるGPUのメモリ容量も変わります。ただし、学習処理と違って、重みを更新する際の活性化関数・損失関数の計算は必要ありませんので、それに係るメモリ容量も必要ありません。そのため、学習と比べると推論時のGPUメモリの利用は重み情報を持ったパラメータのみになるのでメモリ使用量は減ります。

今回は、Cerebras-GPTの7つのモデルを利用しましたが、推論実行にnvidia-smiコマンドで確認したメモリ容量は以下となります。

※注　推論時にはA100 GPUを1枚のみ利用しました。

※Cerebras-GPT 111M, 256M, 590M, 1.3B, 2.7B, 6.7B, 13Bを推論させた際のメモリ使用量

13Bモデルは40,459MiBとなっており、40GBを超えていますので、推論実行にはOut Of Memoryが発生しました。そのため、DeepSpeedを利用し、その結果、13B with DeepSpeedの結果のように5,723MiBまでメモリを削減し実行しました。

以下は推論実行時の結果となります。

実行時には「”There are famous mountain like Mt.Fuji in Japan. Please tell me highest mountain in the world.”」というプロンプトを与え、生成されるまでの時間を測定しています。

なお、今回の測定は、GPUにモデルを常駐させることはせずに実行していますので、GPUへのモデルのロード処理も時間に含めています。

※Cerebras-GPT 111M, 256M, 590M, 1.3B, 2.7B, 6.7B, 13Bの推論実行時間

1.3Bモデルまでは、30秒以内で生成出来ましたが、2.7Bモデルを超えると1分を超えます。13Bモデルに至っては、DeepSpeedでCPU Offloadを使ったため、更に遅くなりました。ただ、大半の時間はモデルをGPUに展開する時間となるため、プログラム次第で推論時間の短縮は可能です。

補足ですが、以下は各モデルの推論(生成結果)です。情報が当たっているかは別として、生成の精度的にはやはり大きなモデルを利用したほうが良いと思います。

最後に

今回、Cerebras-GPTの7つのモデルを利用し、学習及び推論を実行させた際の時間を測定しました。有名なLLMと比べると今回は最大モデルでも130億パラメータと小さいモデルで測定しましたが、この規模でもそれなりのGPU構成が必要になることが分かったと思います。弊社ではCerebras-GPTをリリースしたCerebras社が誇る世界最大チップを搭載したCS-2を販売しています。この製品はLLMを高速に処理できるAIアクセラレータとなっております。

今回までは、Cerebras-GPTに焦点を当てた記事でしたが、今後はCS-2でLLMを実行した際の結果などを記載させていただきたいと思います。

CS-2の記事については、日本ヒューレット・パッカード合同会社様と共同で【Cerebras × Superdome Flex】連載ブログ(全4回)を掲載しています。ご興味がある方はこちらもご覧ください。

第1弾	【GPUではなく、AI専用プロセッサーという選択肢】AIモデル開発にCerebras×Superdome Flex	日本ヒューレット・パッカード合同会社様
第2弾	【巨大ディープラーニングモデルにミッションクリティカルサーバーという選択肢】AIモデル学習にSuperdome Flex× Cerebras	東京エレクトロンデバイス
第3弾	【Cerebras × Superdome Flex第3弾】 Superdome Flexの特徴に迫る！	日本ヒューレット・パッカード合同会社様
第4弾	【Cerebras × Superdome Flex第4弾】Superdome Flex×Cerebras共同検証結果と解説	東京エレクトロンデバイス