こんにちは、初めましての方は初めまして。株式会社 Fusic で機械学習エンジニアをしている瓦です。「2024 年は日記をつけてみよう」と思って書き始め、一日坊主で終わってしまった一ヶ月前の日記を眺めながら、日記を付けることの難しさに絶望しています。
この記事では llama.cpp を使って、AWS Lambda で LLM を動かしてみます。LLM を使おうとすると、大きなメモリの載っている GPU を用意して、という流れに大体なりますが、もちろん誰でも強い GPU が使える環境を得られるとは限りません。また GPU を使おうとするとそれなりの金額を覚悟しないといけないことも多いです。そこで CPU でも動く環境を作って、サーバレスな環境で動くかを試してみようと思います。
準備
モデルの量子化
今回は Stablity.ai が公開している Japanese StableLM-3B-4E1T Instruct を Lambda で動かしてみます。そのままでは Lambda のメモリ上限に引っ掛かって実行できないため、量子化を行いモデルの軽量化を行います。軽量化の方法はいくつかある[1]のですが、今回は CPU を対象にして LLM を動かしたいので、llama.cpp を使いたいと思います。
llama.cpp のビルドは README に従って以下のコマンドを叩けば出来ます。
git clone https://github.com/ggerganov/llama.cpp
cd llama.cpp
make
これで llama.cpp を使う準備が出来たので、モデルの量子化を行います。これも README の Prepare and Quantize に基本的に従えばよいです。ただ、Japanese StableLM-3B-4E1T Instruct は convert.py では量子化できないため、convert.py の代わりに convert-hf-to-gguf.py を使って量子化を行います。
python convert-hf-to-gguf.py <MODEL DIR>
./quantize <MODEL_DIR>/ggml-model-f16.gguf <MODEL_DIR>/ggml-model-Q4_K_M.gguf Q4_K_M
quantize の最後の引数が量子化の方法です。今回は README のまま Q4_K_M を用いました[2]。
これで 4bit への量子化が出来ました。実際に量子化したモデルを ls コマンドで表示してみると、5.3G から 1.6G へとサイズが小さくなっていることが分かります。
-rw-r--r-- 1 ubuntu ubuntu 6.0K Feb 8 09:17 README.md
-rw-r--r-- 1 ubuntu ubuntu 868 Feb 8 09:17 config.json
-rw-r--r-- 1 ubuntu ubuntu 5.2K Feb 8 09:17 configuration_stablelm_epoch.py
-rw-r--r-- 1 ubuntu ubuntu 111 Feb 8 09:17 generation_config.json
-rw-r--r-- 1 ubuntu ubuntu 5.3G Feb 8 09:35 ggml-model-f16.gguf
-rw-r--r-- 1 ubuntu ubuntu 1.6G Feb 8 09:38 ggml-model-q4_K_M.gguf
-rw-r--r-- 1 ubuntu ubuntu 5.3G Feb 8 09:27 model.safetensors
-rw-r--r-- 1 ubuntu ubuntu 28K Feb 8 09:17 modeling_stablelm_epoch.py
-rw-r--r-- 1 ubuntu ubuntu 99 Feb 8 09:17 special_tokens_map.json
-rw-r--r-- 1 ubuntu ubuntu 2.1M Feb 8 09:17 tokenizer.json
-rw-r--r-- 1 ubuntu ubuntu 4.7K Feb 8 09:17 tokenizer_config.json
本当に動くのか気になるので、これも README に書いてある通りにローカルで実行してみます。
./main -m ./models/mymodel/ggml-model-Q4_K_M.gguf -n 128 -p 人工知能とは、 --log-disable
# >> 人工知能とは、自ら学習し自律的に意思決定するコンピュータのことです。将来的には人間の仕事を奪うものとして議論されていますが、「ベイズ統計」と呼ばれる学習方法によってAIは多くの分野で大きな成果を上げつつあります。
ちゃんとプロンプトを入力としてテキストの生成が出来ていそうです。
イメージの用意
イメージベースの AWS Lambda を使用したいので、Dockerfile を用意して ECR にプッシュします。本当はモデルは S3 に置いたりした方がよいと思うのですが、今回は面倒だったのでモデルごとイメージの中に突っ込みます。用意した Dockerfile は以下のようにしています。
FROM python:3.11
RUN apt update && apt install -y git
RUN pip install awslambdaric llama-cpp-python
COPY ggml-model-q4_K_M.gguf ${LAMBDA_TASK_ROOT}/models/
COPY app.py ${LAMBDA_TASK_ROOT}/
ENTRYPOINT [ "/usr/local/bin/python", "-m", "awslambdaric" ]
CMD [ "app.lambda_handler" ]
python 側では llama-cpp-python を使って量子化したモデルを読み込みます。
from llama_cpp import Llama
llm = Llama(model_path="/models/ggml-model-q4_K_M.gguf")
def lambda_handler(event, context):
output = llm(event["text"], max_tokens=event.get("max_tokens", 128))
return output
生成の際に top_p や temperature などのパラメータも使用できます。今回は単に動くかを確かめたかったのでデフォルトのままとしています。
これらのファイルと重みを量子化したファイルを同じディレクトリに配置し、イメージのビルドを行って ECR へと push します。
実行
Lambda を作って実行してみます。コンテナイメージから作成を行い、設定でメモリ上限やタイムアウトの時間を設定します。設定が出来たら、Lambda のテストで実際に実行してみます。
Lambda で動かすことが出来ています！上の画像では「人工知能とは、」に続く文を生成させており、しっかり続く文が生成されていることが分かります。また、メモリは 2GB 未満で時間も二分ほどで実行できていることが分かります。この実行結果はコンテナを立ち上げてモデルを読み込み、テキストを生成するまでの時間の合計なので、おそらく連続で実行するとテキスト生成だけが走る（上の実装では lambda_handler 部分だけ走る）ので、もっと早くテキスト生成出来ると思います。結構現実的なメモリ消費量と時間で実行できているのではないでしょうか？
まとめ
3B という、LLM の中では大きくないサイズのモデルではありますが、llama.cpp を使うことで AWS Lambda で動かすことが出来ました。今回の方法を使えば、独自のデータを使って学習させたモデルを Lambda で動かして自分のサービスに組み込んだりすることが簡単に出来ます。もちろん量子化している分性能は低下しているので、実際のサービスに組み込む場合は間違いが許容出来るかどうかも合わせて検討していくのがいいと思います。
ちなみに、この調子で 7B も動くんじゃないかと思って Swallow-7B-hf でも試してみたのですが、こちらはうまくいきませんでした… ローカルでは出力まで動いており、Lambda 上でもモデルの読み込みまではうまくいっていたので、実際のテキストを入力したときのメモリ消費が厳しいのかなと思っています（「7B も Lambda で動かせたよ」って方がいればこっそり教えてください…）
最後に宣伝になりますが、機械学習でビジネスの成長を加速するために、Fusicの機械学習チームがお手伝いたします。機械学習のPoCから運用まで、すべての場面でサポートした実績があります。もし、困っている方がいましたら、ぜひFusicにご相談ください。お問い合わせから気軽にご連絡いただけますが、TwitterのDMからでも大歓迎です！
Discussion