2024年なので自分でLLMを学習させてみたい。企業が用意しているGPT-4やGeminiみたいなありあわせのものではなく、自分で好きなようにカスタマイズしてみたい。 企業と比べれば個人でできることはたかが知れているが、実際のところ本気を出したらどれくらいの性能のLLMを作成できるだろうか？ モデルの大きさやデータの量よりも、データの質のほうが重要という噂もある。個人レベルでも案外良いものができるのではないか？
LLMを学習させるには
- モデル
- 学習データ
- 計算機資源
の3つを揃える必要がある。また、学習スクリプトも必要だ。
モデル
公開されている様々なOSSモデルがhuggingface hubにある。pretrain済みのものを使うのがよいだろう。最初から学習させるのもやってみたいが、コスト的に厳しい。
いわゆるLocalLLMはMetaが公開しているllama2とその派生がデファクトを取っているが、llama2はライセンスが気に食わない。Mistralという企業がApachev2ライセンスのllama2互換のモデルを公開している。そして、Mistralのllama2互換モデルを日本語データで追加学習させたモデルshisaが公開されている。このshisaに追加学習してカスタマイズができるというのをひとつの目標にしたらよさそうだ。https://huggingface.co/augmxnt/shisa-7b-v1
shisaを公開しているaugmxntがどういう団体なのかはよくわからないが、akavirtualという企業から計算機資源の提供を受けたということが書いてあるので、関係する何かなのだろう。https://akavirtual.com/
学習データ
おそらく一番重要な部分で、検討が必要。とりあえずは公開されているデータでどれくらいのコストがかかるのか見積もりたい。
計算機資源
H100が使えるなら使いたい。しかしH100は500万円以上する。AWSなどのクラウドサービスを使うのも検討したが、長く何回も学習させるにはコストが高すぎる。
最新のGPUは無理、クラウドも避けたい。学習に向いた安いGPUを検討すると、現状RTX3090の中古を買うのが最も優れていそうだ。学習に最も必要なのはGPUのメモリで、RTX3090はVRAMが24GBある。ヤフオクでRTX3090の中古は10万円前後で取引されている。A100の80GBには及ばないし、3000番台で新しいGPUではないが、この価格で24GBあるのはRTX3090だけだ。同じく24GBあるRTX4090は中古でも30万円くらいする。
とりあえずRTX3090を2台揃えた。本当は4台にしたいが、マザーボードを変えなければならないので2台で様子見する。
学習スクリプト
Microsoftが公開しているOSSのDeepSpeedを使うとZeROという仕組みで少ない計算機資源で大きなモデルの学習ができる。
ZeROのアイデアは単純だ。これまで各GPUで重複してそれぞれが持っていたパラメータを代わりに分割して部分部分をそれぞれのGPUで持っておき、必要なタイミングで必要な場所に集める。これで必要なメモリ量は単純には1/（GPUの台数）倍になる。データの移動と計算をオーバーラップすることでそれらのうちの速い方にかかる時間は遅い方の時間に隠れて見えなくなる。
ZeROはstageが0, 1, 2, 3と4種類あるが、学習にかかるどのパラメータを優先順にどこまで分散するかで分かれている。僕の手元にある計算機は足りないので当然stage3を使う（つまりOptimizerのパラメータ、gradientデータ、モデルのパラメータ全部を分散して持つ）。それでも足りないので、ZeROとは別で用意されているCPU offloadingの機能も使う。これはGPUメモリの足りない分をCPUメモリで補う機能だ。
そしてCPUメモリも足りないので、NVMeストレージを買ってきてswap領域にして誤魔化している。
データの移動がボトルネックになりそうだが、こういう貧者の工夫は楽しい。
huggingfaceが公開しているtrlというOSSのexampleの学習用スクリプトを少し改造してやるとDeepSpeed ZeRO stage 3 with CPU offloadingで学習ができた。
