大規模言語モデルの学習に必要な計算量を試算する
はじめに
最近は大規模言語モデルを作っています｡
10B級のモデルを作れるリソースを使わせていただける予定なのですが､具体的なFLOPの推定が必要そうだったので､軽く試算してみることにしました｡
計算が間違っていたらすみません｡
試算サイトを活用する
以下のサイトを使うと､モデルサイズ・学習トークン数をもとに､学習に必要な計算量FLOPsやestimated lossを推定できます｡素晴らしいですね｡
(FLOPsとFLOPSの使い分けについてはこちら)
まずは､このサイトの信頼性を確認する必要があります｡
実測
最近､0.35 Bモデルの学習を行ったので､実測値と比べてみます｡
学習時に出力されるLogや､データセットのトークン数(2.8B)などをゴニョゴニョ計算すると､FLOPsなどを推定できます｡
→ 2 * 10^18 FLOPsでした｡
後の計算のため､1 GPUが1時間あたりに計算できるFLOPSも計算しておきます｡今回の学習には､A100 (80GB)x2を用い､4.8時間を要したので､それらで割ると､
4.5 * 10^17 FLOPs/GPU/hour となりました｡
推定値
これに対し､上記のサイトで0.35 Bモデルで2.8B tokenを学習させたと仮定したときの推定計算量は 5.9 * 10^18 FLOPsでした｡
3倍程度､推定値を過大評価してますが､目安としては使えそうです｡
ちなみに､予測誤差に対応するestimation lossは実測が3.0､予測が2.5でした｡この値も､目安としては使えそうです｡
10Bモデルの計算量を試算する
10Bモデルで600B tokenを学習させるための試算FLOPsは3.6*10^22でした｡
ざっくり､これをA100の計算速度で割ると､79535 GPU・hrが得られました｡
仮に20 GPUで計算を回すとすると､約4000時間 (166日)を要するようです｡
更にざっくり､H100がA100の3倍の性能を持つと仮定するとでは､55日で計算が終わります｡
さらに､当該サイトの試算値が2-3倍程度､過大評価されていそうなことも踏まえると､値をとりあえず2で割って※､30日くらいで計算が終わるような気がしました｡
※並列化をすればするほど､通信のボトルネックなどの影響が生じやすくなり､並列化効率が落ちたりします｡このあたりの要素も､正確な試算には踏まえる必要があります｡ また､H100だとTransformerEngineという高速化手法を使えたりもするようです｡
50Bモデルの計算量を試算する
ついでに､この作業を50Bモデルで行ってみます｡
50Bモデルで600B tokenを学習させるための試算FLOPsは1.8*10^23でした｡
Lossの低下は0.1程度のようですが､体感では､10bと数十bの間には大きな性能差があるように感じます｡
以下､H100換算で､推定値を2で割ったときの計算時間を示します｡
20 GPU: 138 days
50 GPU: 55 days
100 GPU: 28 days
200 GPU: 14 days
まとめ
大規模言語モデルの学習に必要なGPUリソースを軽く試算しました｡
どれくらいのリソース投入が必要なのかを予測するサイトが非常に役立つことがわかりました｡
感想
やはり､計算リソースの確保は大事な課題だと思いました｡
しれっと200 GPUなどと書いてますが､1枚あたり550万円はするので､GPU代だけでも11億円かかるようです｡
3/25追記: もう少し正確な試算
FLOPSが､学習環境やモデルサイズによって､かなり変わることが分かってきました｡
以下は､約30 b tokenのデータを0.1, 2.7bのモデルで学習させた際の結果です(A100 x 8)｡
モデルサイズが小さいほど､flopsが落ちている理由は明確ではありません｡
(モデルが小さいほど､バッチサイズが大きくなり､データ通信等で遅延が生じている感じでしょうか｡)
ざっくり､モデルの計算時間は､以下の式で推定できる模様です｡
(基準となるモデルの計算時間) x (モデルサイズの比) x (トークン数の比) x (FLOPsの比)
今回の場合は､
(0.1bモデルの計算時間: 0.64) x (2.7 / 0.136) x (1/1) x (63/140)≒5.6 days
となり､実測(5.5 days)とほぼ一致します｡
基本的に､計算量はモデルサイズとトークン数に比例するので､いかに計算効率を落とさずに､GPUで並列化するか､という(当たり前の)に帰結する模様です｡
4/2追記
こちらの資料(第四回､p104~)にも推定方法がありました｡
基本的な考え方は上と同じですが､
(推定時間) = (基準となる条件での学習時間) x (データ数の比) x (モデル数の比) x (GPU枚数の比) x (補正係数)
で出せそうです｡
補正係数は､GPUの種類(例えばH100はA100の3倍程度の性能?)や､分散学習時のFLOPs効率(通信などがボトルネックになりやすい) などで決まります｡
資料では､175bモデルを300bトークン､A100を1000枚で学習させたところ､36日を要したとの記載がありました｡
それを､今回想定するケースに換算したのが､以下の表になります*｡
*今回のコンペでは､H100が30日 x 18枚相当､提供される予定です｡
結局のところ､H100を使った時のFLOPsがどれくらいになるのか､更にはtransformerengineのような最適化ツールを使ったときに､どれくらいの加速効果があるのか(→factor)がわからないと､正しい推定値を出すのは難しそうです｡
ファインチューニングや予備期間を含めると､事前学習に使えるのは20-30日程度になりそうです｡
以上の結果をまとめると､事前学習に使えそうなtoken数は 200-600 bとなりました｡ 幅がありすぎて､ちょっと困っているところです｡もう少し正確に推定したいところです｡良い計算法をご存知の方がいたら､押して頂けると助かります｡
この記事が気に入ったらサポートをしてみませんか？