【LLM】モデルの重みをWebLLM形式に変換する方法
LLMをブラウザで推論させる方法にwebLLMを使う方法がある。webLLMを使うとWASMとWebGPUを使用して量子化モデルをブラウザで動かすことが出来る。
WebLLM は、ハードウェアアクセラレーションを使用して言語モデル推論を Web ブラウザーに直接提供する、高性能なブラウザー内 LLM 推論エンジンです。すべてがサーバーのサポートなしでブラウザー内で実行され、WebGPU で高速化されます。
しかし、web-llmではwebGPU用の変換をする必要があるので、ここではモデルをwebGPU用にMLC-LLMに変換する方法を紹介する。
やったこと
独自でfine-tuning したモデルAtotti/TinySwallow-GRPO-TMethod-experimental
をq4f32_1
で量子化したMLC形式のモデルAtotti/TinySwallow-GRPO-TMethod-experimental-q4f32_1-MLC
に変換した。
ブラウザで推論するアプリケーションの開発も行った。それはまた次の記事で紹介する。
環境構築
Dockerを使いましょう。Dockerさえ入ってしまえばLinuxのOSやCUDAのバージョンが関係ないのでお勧めです。
DockerでCUDA GPU環境を作る方法は、以下の記事などを参照してください。
変換
まず Docker イメージをビルドする。
FROM nvidia/cuda:12.3.2-cudnn9-devel-ubuntu22.04
ARG PYTHON_VERSION=3.11
ENV DEBIAN_FRONTEND=noninteractive
ENV HOME /app
WORKDIR $HOME
RUN apt-get -y update && apt-get -y upgrade && \
apt-get install -y --no-install-recommends \
git \
build-essential \
libssl-dev \
zlib1g-dev \
libbz2-dev \
libreadline-dev \
libsqlite3-dev \
liblzma-dev \
liblzma-dev \
libffi-dev \
curl \
clang \
git-lfs \
make \
pkg-config \
libgoogle-perftools-dev
COPY /uv /uvx /bin/
RUN uv python install $PYTHON_VERSION
ENV HF_HUB_CACHE /app/.cache/huggingface
RUN curl https://sh.rustup.rs -sSf | sh -s -- -y
ENV PATH /root/.cargo/bin:$PATH
######################### Install CMake 3.27.4 #########################
WORKDIR /root/temp
RUN curl -OL https://github.com/Kitware/CMake/releases/download/v3.27.4/cmake-3.27.4.tar.gz
RUN tar -xzvf cmake-3.27.4.tar.gz
WORKDIR /root/temp/cmake-3.27.4
RUN ./bootstrap -- -DCMAKE_BUILD_TYPE:STRING=Release
RUN make -j4
RUN make install
WORKDIR /root
RUN rm -rf temp
CMD ["bash"]
services:
hf2webllm:
container_name: hf2webllm
image: atotti/hf2webllm:latest
working_dir: /app
build:
context: .
dockerfile: ./Dockerfile
deploy:
resources:
reservations:
devices:
- driver: nvidia
count: 1
capabilities: [gpu]
ビルドして起動
docker-compose build
docker-compose run --rm -it hf2webllm bash
以降のコマンドはコンテナ内で行う。ビルドが完了するとbashが立ち上がるようにしてあるので、そのシェルを使えばよい。(個人的にはコンテナにVSCodeをアタッチして使うのが好み)
仮想環境に入る
uv init .
uv venv
source ./.venv/bin/activate
MLC-LLMのビルドを行う。wheelからインストールもできるが、執筆時現在、配布されているwheelにはバグがあったのでコードからビルドする。
# clone from GitHub
git clone --recursive https://github.com/mlc-ai/mlc-llm.git && cd mlc-llm/
# create build directory
mkdir -p build && cd build
# generate build configuration
python ../cmake/gen_cmake_config.py
# build mlc_llm libraries
cmake .. && cmake --build . --parallel $(nproc) && cd ..
configはCUDAだけyesにしてそれ以外はnoで良いはず。
ビルドしたMLC-LLMをインストールする
cd /app
uv pip install -e ./mlc-llm/python/
TVMのインストール
cd /app
uv add --pre -U -f https://mlc.ai/wheels mlc-ai-nightly-cu123
mlc-llmとTVMのインストールが正常に出来ているか確認する
python -c "import mlc_llm; print(mlc_llm)"
python -c "import tvm; print(tvm.__file__)"
以下の様な出力が得られればOK
(app) root@398dd1afee5c:~# python -c "import mlc_llm; print(mlc_llm)"
<module 'mlc_llm' from '/app/.venv/lib/python3.11/site-packages/mlc_llm/__init__.py'>
(app) root@398dd1afee5c:~# python -c "import tvm; print(tvm.__file__)"
/app/.venv/lib/python3.11/site-packages/tvm/__init__.py
まず、HuggingFaceから変換したいモデルをクローンする。
mkdir models && cd models
git lfs install
git clone https://huggingface.co/Atotti/TinySwallow-GRPO-TMethod-experimental
cd ..
次に、重みの変換を行う。量子化の種類をq4f32_1
と設定している。利用可能な量子化オプションを確認できる。
mlc_llm convert_weight ./models/TinySwallow-GRPO-TMethod-experimental/ \
--quantization q4f32_1 \
-o dist/TinySwallow-GRPO-TMethod-experimental-q4f32_1-MLC
次に、チャットテンプレート等のコンフィグを作成する。--conv-template qwen2
はそのモデルに合ったものを設定する必要がある。TinySwallowはQwen2ベースであるのでQwen2を指定している。指定可能なテンプレートを確認できる。ここに無い場合は自分でカスタムする必要がある。
mlc_llm gen_config ./models/TinySwallow-GRPO-TMethod-experimental/ \
--quantization q4f32_1 --conv-template qwen2 \
-o dist/TinySwallow-GRPO-TMethod-experimental-q4f32_1-MLC/
ここまでで変換は完了している。あとはdist/に出力されたモデルをHuggingFaceにアップロードするなりして利用出来る。
確認
実行して正しく変換出来ているかを確認する。
mlc_llm chat dist/TinySwallow-GRPO-TMethod-experimental-q4f32_1-MLC/
変換後のモデルで正しく動作することが確認できた。
>>> こんにちは
こんにちは! 👋
何かお手伝いできることはありますか? 😊
ハマった箇所
重みの変換mlc_llm convert_weight ./models/TinySwallow-GRPO-TMethod-experimental/ \ --quantization q4f32_1 \ -o dist/TinySwallow-GRPO-TMethod-experimental-q4f32_1-MLC
を実行したときに以下の様なエラーが発生した。
TypeError: PagedKVCache.attention_with_fused_qkv() missing 1 required positional argument: 'sm_scale'
当該箇所をライブラリのGitHubを見に行くと、以下の様になっており、PagedKVCache.attention_with_fused_qkv() missing 1 required positional argument: 'sm_scale'
というエラーは発生し得ないように思った。
paged_kv_cache.attention_with_fused_qkv(
layer_id, qkv, self.num_attention_heads, sm_scale=self.head_dim**-0.5
)
ところが、手元の環境の.venv
の中の当該コードを見に行くと、以下のように sm_scale
が省かれておりバグっていた。
paged_kv_cache.attention_with_fused_qkv(layer_id, qkv, self.num_attention_heads)
インストール時にwheelからインストールしたために、最新版と配布されているwheelの差分が問題だったので、GitHubから直接ビルドし解決した。
※確認したところ、この修正は執筆時点で30分前にマージされた修正だったので仕方ない...
おわりに
次の記事では変換したMLC形式のモデルを使ってブラウザ上でLLMを実行する方法を紹介する。(こっちの方がハマりどころが多かった)
さらに、この方法でローカルLLMを活用したWebアプリケーションについても作成したので紹介したい。
また、いずれTinySwallow-1.5B-Instructに強化学習(GRPO)でfine-tuningし、reasoningと出力のフォーマット指示追従性能が(おそらく)向上した今回変換したモデルの学習についても紹介したい。
Discussion