【Petals】3人で力を合わせてLlama-2-70bを動かす【Colab】

はじめに

今回の目標はこの巨大モデルLlama-2-70bを動かすことです。

700億パラメータということで、4bit量子化をかけても40GB程度のVRAMが要求されます。残念ながらColabの無料枠には乗りきりません。

今回はPetalsというOSSを使用することで、Colab上での実行を試みます。

Petalsとは

簡単に説明すると、モデルをいくつかに分割し、複数のマシンでそれぞれの演算を担うことで、大規模なモデルでの推論を実現するというものです。

現在ホストされているモデルは、こちらのHealth Monitorで確認することができ、誰でも利用することができます。

サポートされているモデルとしては、BLOOM, LLaMA, Falconなどが挙げられ、どれもOSSとしては最大規模のパラメータを有します。

ホストされているモデルで推論してみる

まずはPetalsでホストされているモデルを用いて推論ができる様子を確かめます。

環境はGoogle Colabで、ラインタイムはCPUで大丈夫です。

まずはpetalsをインストールします。

!pip install git+https://github.com/bigscience-workshop/petals

次に使用するモデルを指定します。ここではPetalsで多くホストされているpetals-team/StableBeluga2を使用します。

import torch
from transformers import AutoTokenizer
from petals import AutoDistributedModelForCausalLM

model_name = "petals-team/StableBeluga2"

tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False, add_bos_token=False)
model = AutoDistributedModelForCausalLM.from_pretrained(model_name)

最後に、こちらのモデルを用いて推論を行います。

inputs = tokenizer('A cat in French is "', return_tensors="pt")["input_ids"]
outputs = model.generate(inputs, max_new_tokens=3)
print(tokenizer.decode(outputs[0]))

初回のみ少し時間がかかります。2回目以降では、私の環境だと5秒程度で生成されました。

出力結果

A cat in French is "un chat"

うまく推論できていますね。

モデルをホストしてみる

では次に、自分でモデルをホストしてみましょう。
まずは軽量モデルとして、Llama-2-7bを試します。

先ほどと同様、まずはpetalsをインストールします。ラインタイムはT4 GPUを使用します。

!pip install git+https://github.com/bigscience-workshop/petals

次にLLaMAモデルを使用するために、HuggingFaceにログインします。

PetalsでLLaMaモデルを使用するためには、Metaへの利用申請およびHuggingFaceのアカウント登録が必要です。申請はこちらから行えます。

アクセストークンの入力が要求されるので、HuggingFaceで取得したトークンを入力してください。

!huggingface-cli login

最後にサーバーを走らせ、モデルをホストします。

!python -m petals.cli.run_server meta-llama/Llama-2-7b-hf

しばらくモデルのロードなどが走ります。「Started」と表示されたら準備は完了です。

実際にホストされているかをHealth Monitorで確かめましょう。

ちゃんとホストされていることがわかりますね。では別のNotebookを開いて、先ほどと同様にLlama-2-7bで推論を行ってみましょう。ここではAIを活用する上での注意点を教えてもらいました。

入力プロンプト

'The points to keep in mind when using AI are　'

出力結果

The points to keep in mind when using AI are　1) to make it a tool, not a goal, 2) to be aware of the limitations of AI, and 3) to keep in mind the humanity of those who are affected by AI.

結構いい感じのことを返してくれました。うまく推論できていますね。
生成時間は20秒くらい（1トークンあたり0.4秒くらい）で、速度もそこそこ出ています。

Llama-2-70bで推論してみる

では最後に本題である、Llama-2-70bの実行を試みます。

まずは先ほどと同様にLlama-2-70bをホストし、Health Monitorでモデルの状態を確かめます。

一部のブロックは実行されているものの、計算資源が不足しているようで、全てのブロックは実行できていません。
当たり前ですが、このままでは推論することができません。

そこで有志2人の力を借り、同様にLlama-2-70bをホストしてもらいます。
今一度モデルの状態を確かめてみましょう。

これで全てのブロックを実行することができました。

最後に、実際に推論を実行してみます。

入力プロンプト

'The points to keep in mind when using AI are　'

出力結果

The points to keep in mind when using AI are　as follows:
・Do not use the AI to make the final decision.
・The AI does not have a sense of humanity.
・Do not use the AI to replace human beings.
・Do not use

うまく推論することができました！
生成時間は2分半くらい（1トークンあたり3秒くらい）で早くはないですが、70bを自分で動かせるのは大きいですね。

おわりに

LLMはその計算負荷の大きさから、なかなか個人が手を出しにくい側面もあります。
このように多くの人が計算資源を出し合うことで、大規模モデルを実行するという取り組みは面白かったです。

興味を持たれた方はぜひ試してみてください！