今春、8GBの環境で生成AIがどこまで動くのかを模索していた*1。パフォーマンスを気にしなければ、GPTだろうがDiffusionだろうが、GPUの無いRAM=8GBの環境、PCでもAndroidスマートフォンでも動作はするのだ。ただし推論であれば、という条件付きだった。つまり、誰かが提供してくれる小規模なモデルを用いる前提であり、そのモデルに独自の情報を付与する強化学習に対応することは現実的ではなかった。
そのような現状に微かな希望をもたらしてくれたのが、llama2.cだった。
llama-2アーキテクチャのシンプル、かつミニマムな推論エンジンを開発するという、Andrej Karpathy氏による週末のクリエイティブな活動は、学習モデルの小規模さから、推論だけでなく強化学習までもRAM=8GBの環境で実行できるのだ。
実用的なところは置いて、モデルの生成から、そのモデルを用いた推論まで、一連の流れをRAM=8GBの環境で追体験してみた。この投稿は、その手順と記録だ。
なお推論の実行については、LinuxでもWindowsでも簡単に実行できる。しかしモデルの生成について、Windowsでの実行は難しい。以下のような障害がある。
- プログラムから呼び出されるシェル・スクリプトの書き直し
- 学習データ・テキストのエンコード対応
- sentencepieceのWindows用ビルド
- 学習データのコンパイル→torch.compileがWindowsに対応していない
Windowsについては、推論エンジンの初回実行のみ取り上げる。
前提
この投稿では、次の環境を用いている。
| Windows | Windows 11 Pro 22H2 Visual Studio 2022 17.6.5 |
| Linux | Clear Linux 39850 Python 3.11.4 |
この投稿では、次のフォルダ構成を前提としている。
| 作業フォルダ | ~/work |
| llama2.cフォルダ | ~/work/llama2.c |
| Python仮想環境 | ~/work/llama2.c/myenv |
| sentencepieceフォルダ | ~/work/sentencepiece |
推論エンジンのビルド、実行
まず作業フォルダに移動し、llama2.cのレポジトリをローカル環境に保存する。
cd ~/work git clone https://github.com/karpathy/llama2.c.git
生成されたllama2.cフォルダに、学習済みモデル"stories15M.bin"をダウンロードする。
Windows
実行ファイルを生成するためにはVisual Studio、あるいはC++ビルドツールをインストールしておく必要がある。インストール方法は過去の投稿で紹介している。Windowsを対象とした、いずれかのページを参照してほしい。
impsbl.hatenablog.jp
環境が整ったら、"Developer PowerShell"から次のコマンドを実行する。
.\build_msvc.bat .\run.exe .\stories15M.bin .\run.exe .\stories15M.bin -i "In Japan, there is a tradition of dancing in circles during summer nights, "
実行環境のCPUは第6世代Intel Core i7だ。推論のパフォーマンスは300トークン／秒程度なのが分かる。物語の生成はランダムのようではあるが、インストラクションを与えた場合には、それに関連するような物語を生成していることも分かる。
 |
 |
|
Linux
この投稿での実行環境はClear Linux*2だ。しかしディストリビューションに依存する設定、作業はないため、他のディストリビューションでも共通に機能するはずだ。
Windows同様、実行ファイルを生成し推論を実行する。"runomp"を指定しているのは、並列処理に対応するためだ。
make runomp ./run ./stories15M.bin ./run ./stories15M.bin -i "In Japan, there is a tradition of dancing in circles during summer nights, "
Python環境の準備、sentencepieceの導入
ここからの話題はLinuxのみを対象とする。各種Pythonスクリプトの実行、並びにsentencepiece*3を導入するための環境を作成する。
sentencepieceはデータの前処理に用いられる。その実行ファイル、ライブラリは次のフォルダへ格納される。
| 実行ファイル | /usr/local/bin |
| ライブラリ | /usr/local/lib64 |
環境によっては、ライブラリ・フォルダのパスを通しておく必要がある。また環境変数として設定を永続化するために、「~/.config/environment.d/envvars.conf」*4にも追記しておくとよい。
# Python環境 cd ~/work/llama2.c python -m venv myenv source myenv/bin/activate python -m pip install --upgrade pip pip install -r requirements.txt # sentencepieceの導入 cd ~/work git clone https://github.com/google/sentencepiece.git cd sentencepiece mkdir build cd build cmake .. make -j $(nproc) sudo make install sudo ldconfig -v #ライブラリ・パスの登録 export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib64
機械学習：前処理
GitHubの「training」で示されているコマンドを実行することによって、推論エンジンの実行でダウンロードしたstories15M.binを生成できるのだという。しかし、その生成にはクラウド環境のGPU4台で並列処理して、1日を要している。そこで「custom tokenizer」で示されているコマンドを実行する。
第7世代Intel Core i5で一連の作業を実施したところ、30分ほどを要した。
cd ~/work/llama2.c python tinystories.py download python tinystories.py train_vocab --vocab_size=4096 python tinystories.py pretokenize --vocab_size=4096
「training」のコマンドでは語彙の大きさ（vocab size）として32000が指定されていた。「custom tokenizer」では、代わりに4096が指定されている。語彙サイズが大きいほど、モデルの精度は向上するものの、用途によってはオーバースペックな場合もある。適切な精度が期待できる、限られた語彙サイズを追求し、指定することで、モデルはより小さくなり、結果として推論処理が高速になる。
”python tinystories.py download”を実行すると、”TinyStories_all_data.tar.gz”がダウンロードされる。該当コマンドを実行せず、”tinystories.py”に記述されたURLから直接ダウンロードしても良い。その場合、件のファイルを展開し、収録されている全てのjsonファイルを次のフォルダへ収納しておく必要がある。
~/work/llama2.c/data/TinyStories_all_data
機械学習：モデルの生成
スワップの生成
モデルを生成する前に、機械学習環境を整えておく。まずスワップを設定する。RAM=8GBの環境では、モデル生成中にRAMが不足し、プロセスが強制終了されてしまう。それを回避するために、8GBのスワップを生成しておく。
sudo dd if=/dev/zero of=~/work/llama2.c/swap.img bs=1M count=8096 sudo chown 0:0 ~/work/llama2.c/swap.img sudo chmod 600 ~/work/llama2.c/swap.img sudo mkswap ~/work/llama2.c/swap.img sudo swapon ~/work/llama2.c/swap.img
train.pyの編集
機械学習処理を担う「train.py」には、多数のパラメータが用意されている。パラメータの調整は、動作不良や実行エラーに通じるリスクを伴う。この投稿では、実用面は犠牲にして、モデル生成処理のループ回数と時間が現実的なものとなるよう、次のようにパラメータを編集している。
| パラメータ | 規定値 | 設定値 | |
|---|---|---|---|
| batch_size | 128 | 8 | |
| compile | True | False | |
| device | cuda | cpu | |
| drop_out | 0.0 | 0.1 | |
| eval_interval | 2000 | 20 | |
| learning_rate | 5e-4 | 4e-4 | |
| max_iters | 100000 | 1000 | |
| warmup_iters | 1000 | 10 |
「compile」をTrueにすることで、PyTorchがモデルをコンパイルする。これによって、より高速なモデルを生成することができる。しかしtorch.compileはPython 3.11に対応していない。この問題を回避するため、ここではFalseとしている。
「device」を"cpu"とすることで、GPUのない環境でもモデルを生成できるようにした。
「eval_interval」「max_iters」「warup_iters」は、モデル生成のループ処理制御とモデル生成に関連している。具体的には、次のループ処理だ。
while True: # ～略 # evaluate the loss on train/val sets and write checkpoints if iter_num % eval_interval == 0 and master_process: # ～略 if losses["val"] < best_val_loss or always_save_checkpoint: # ～略 if iter_num > 0: # ～略 model_export(raw_model, os.path.join(out_dir, "model.bin"), version=0)
デフォルトでは、「max_iters=100000」回の処理をすることになる。第7世代Intel Core i5で処理すると、1回の処理につき、約4.5秒を要する。つまり10万回を終えるには、5日程度を要することになる。これを1時間半程度にとどめるべく、「max_iters=1000」とした。
こうするとデフォルトの「eval_interval=2000」では、if文のネストを全回避することになる。そこで20回に一度、if文のネストを辿れるよう「eval_interval=20」とした。
さらに次の処理で0除算を回避するため、「warup_iters=10」としている。
decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
その他のパラメータは、強化学習の特性に影響を与える部分だ。次の引用で言及されている様に、論文に掲載されているテーブルを参照して調整する必要がある。
Look at the table at the very end of the Chinchilla paper to get a sense of how the Transformer parameters (dim, n_layers, n_heads) grow or shrink together.
📁the table at the very end of the Chinchilla paper
強化学習
ここまでの編集を終えたら、次のコマンドを実行する。第7世代Intel Core i5での実行時間は約90分だった。実行中のhtopを確認すると、一見RAM=8GBでも余裕がありそうに見えるが、前述の如くkilledを経験しており、処理のどこかでRAM容量が逼迫するタイミングがある。
python train.py --vocab_source=custom --vocab_size=4096
 |
 |
この投稿では実行しないが、”train.py”は過去のトレーニングを引き継いで、再開実行することができる。例えば10万回のループ、5日間分の処理だとしても、チェックポイントの出力ごとに処理を中断し、最新のチェックポイント以後のループから再開する、といった対応ができる。
そのためにはパラメータを次のように編集する必要がある。チェックポイントとモデルは再開後の追記に再利用される。
| 初回実行 | 再開実行 | |
|---|---|---|
| init_from | scratch | resume |
処理が完了すると、「~/work/llama2.c/out」に、モデルとチェックポイントが生成される。
 |
 |
"stories15M.bin"のサイズ57.9MBに対して、新しく生成されたモデルのサイズは28.7MBだ。このモデルを用いて実行すると、単位時間当たりの処理トークン数が400強／秒と、パフォーマンスが向上していることが分かる。しかし同時に、英語の文章が稚拙になっており、指定したインストラクションも機能していないに等しいことが確認できる。
全てはモデル生成の学習回数次第なのだ。
python tokenizer.py --tokenizer-model=data/tok4096.model ./run out/model.bin -z data/tok4096.bin
余談：独自データを用いたモデル生成、JSONファイルの構成
"stories15M.bin"はショートストーリー的な140MBのテキスト、50ファイル分を学習したモデルだ。もし、このテキストがカスタマー・サポートの問答だったらどうだろう*5。あるいはローマ五賢帝の一人、マルクス・アウレリウス・アントニウスの自問自答をまとめた『自省録』だったら。汎用的な対話形式を前提としなくとも、何か使い道のありそうなモデルが生成できるのではないだろうか。
”TinyStories_all_data.tar.gz”に収録されているjsonファイルは、1行のテキストで構成されている。その1レコードは、次のような構成だ。
{ "story":"", "instruction":{ "prompt": "", "words": [], "features": [] }, "summary": "", "source": "" }
📁オリジナル・データ
{ "story": "Once upon a time, in a small house, there lived a boy named Will. Will was a selfish boy who never shared his toys with other kids. One day, his mom told him, \"Will, you need to share your toys with your friends.\"\nWill didn't want to share, but he heard his mom and thought about it. The next day, he went to the park to play. He saw a girl named Sue who was sad because she had no toys to play with. Will thought of what his mom said and decided to share his toys with Sue.\nSue was very happy and said, \"Thank you, Will! You are a good friend.\" Will felt happy too, and from that day on, he was not selfish anymore. He always shared his toys with his friends and they all played happily together.", "instruction": { "prompt:": "Write a short story (3-5 paragraphs) which only uses very simple words that a 3 year old child would understand. The story should use the verb \"hear\", the noun \"will\" and the adjective \"selfish\". The story has the following features: the story should contain at least one dialogue. Remember to only use simple words!", "words": ["hear", "will", "selfish"], "features": ["Dialogue"] }, "summary": "A selfish boy named Will learns to share his toys after his mom tells him to do so, and he makes a new friend at the park by sharing his toys with her.", "source": "GPT-4" },
つまり、このような構造の独自データを強化学習に用いることで、独自のデータに基づいたモデルを生成することができる。
今回の投稿で紹介した程度には、RAM=8GBの環境でも強化学習を体験することができる。ChatGPTをはじめとする生成AIサービスが対話形式だからと言って、言語モデルを用いた全てのサービスがそうである必要はないし、そのようなサービスのように汎用的である必要もない。そのようなことを考えると、llama2.cは結構な夢を見させてくれる存在だと思うのだ。
