本記事では、FlexGenをDockerで使う方法を紹介します。
※nvidia-docker2が導入済みの、DockerからGPUが使用できる環境でなければ実行できませんのでご注意ください。
FlexGenとは
ライセンス：Apache-2.0 license
FlexGenとは、自前のPCなど計算資源の小さいマシン上で、大規模言語モデル（Large Language Model: LLM）を動作させることを可能にする技術です。
FlexGenについては、既に概要を詳しく説明している記事が沢山出ていますので、ここでは説明を省略します。詳細は、リポジトリや論文で確認することができます。
論文：High throughput generative inference of large language models with a single gpu
DockerでFlexGenを動作させる
以下のDockerfileを作成します。
FROM ubuntu:20.04
RUN apt-get update && apt-get upgrade -y
RUN apt-get install -y git python3 python3-pip
RUN git clone https://github.com/Ying1123/FlexGen.git -b ablation
RUN cd FlexGen/ && pip install -e .Dockerfileのあるディレクトリで、以下のコマンドを実行します。
docker build -t ubuntu:FlexGen .ビルドが済んだら、コンテナを起動します。起動コマンドは以下の通りです。
docker run -it --gpus all ubuntu:FlexGenコンテナの中のターミナルで以下を実行します。今回は、6.7billionのパラメータサイズの大規模言語モデルを実行してみます。
python3 apps/chatbot.py --model facebook/opt-6.7b無事起動すると、以下のような画面になります。
Humanと記載されているところに、文章を入力すと、応答文がAssistantの欄に返ってきます。
終了するときは、何も入力せずにEnterを押下します。
さいごに
本記事では、DockerでFlexGenを動作させる方法を紹介しました。まだ新しく登場したばかりなので、今後、機能が追加されていくことが予想されます。その時は、また内容を追加しようと思います。
