🧠

70Bモデルのホスティング

ken

2023/12/28に公開

こんにちは、@ken11です。
今日はSpiral.AI株式会社のエンジニアブログとして、70BのLLMをホスティングした話をご紹介したいと思います。

70B

春先に「とりあえず6B、話はそれからだ」と言った人がいたような気もしますが、たしかに最近のLLMはパラメータ数が7B/13B/70Bというのが多いのではないでしょうか。
そうですね、Llama2をファインチューニングしたものも多いので、余計にそういったパラメータ数を見かけることが多くなった気がします。

さて、みんな大好きOpenAIのGPT-3.5が175BだとかGPT-4はおそらく500B~だとかいろいろ言われており、Llama2最大サイズの70Bでも遠く及ばないわけですが、それでも70Bは市場に出回っている(誰でも使える)モデルの大きさとしては現状最大級となっているわけです。
とはいえ実際に使おうと思うと、7Bですら結構な推論コストがかかり、70Bとなるといよいよ莫大になってきます。
そういった事情から逆にパラメータ数を圧縮して性能を保つ方法を模索している企業も最近はよく見かけますね。そうです、どんなにいいモデルでもホスティングして実際に使えなければ意味がないのです。

というわけで、今回は70Bのモデルをどのようにホスティングしたらよいか、という話です。
Spiral.AIでは現在70Bのモデルを実際にホスティングしています。(非公開)

推論にかかるコスト

今回はわかりやすくLlama2の70Bを前提として話を進めたいと思います。
このLlama2 70Bモデルですが、素直にfp16とかで推論する形でロードしようと思うと、だいたい130GB程度のGPUメモリが必要と言われています。load_in_8bit=True　でその半分程度となってきます。
これをいまのAWSの事情に変換すると以下のような状態です。

インスタンスタイプ	インスタンスが積んでいるGPU	1時間あたりの料金	1ヶ月(732時間)稼働した場合のおおよその料金	備考
g5.12xlarge	A10G(24GB) x 4	8.22609$	85万円	load_in_8bit
g5.48xlarge	A10G(24GB) x 8	23.62246$	244万円
p4d.24xlarge	A100(40GB) x 8	44.92215$	463万円
inf2.48xlarge	Inferentia2(32GB) x 12	19.4719$	201万円

1台でこの価格なので、常時稼働させるコストがいかに高いかがよくわかると思います。
これにはAWSの現在抱えるGPUの問題もあり、たとえば上述のインスタンスタイプだとp4dというのは明らかに学習向けであり推論用の選択肢にはなりづらいわけです。
A100 x 2のような柔軟な選択肢はなく、この点ではアリババクラウドの方が有利と言わざるを得ません。
また、AWSでは p4de というA100(80GB)のインスタンスファミリーや、p5e(H200) g6(L4) g6e(L40S) というHopper世代を搭載したファミリーがローンチを控えており、この辺が来ることで学習・推論の選択肢は増える予定ですが、現状g5系なども枯渇していたりする状況なので、Hopper世代が来たとして入手性が高いかどうかは疑問です。

このような状況下なので、そもそも70Bをホスティングするというのは結構難易度が高いというわけです。

サーバレス推論という選択

そこで今回僕が選択したのが、cerebriumというサーバレスGPUのサービスです。

https://www.cerebrium.ai/

cerebirumの特徴

cerebrium最大の特徴はサーバレスGPUという点です。機械学習モデルの推論に特化したサービスで、推論用に書かれたペライチのコードをプッシュするだけで、AWS Lambdaのような使い心地で簡単に推論処理を実行することができます。

提供されているGPUにはA6000/A5000などAWSにはない選択肢も用意されており、A100以外は比較的潤沢(にみえる)^[1]状況というのが個人的に気に入っている点です。

料金はそれなりなので、もちろん使う時間が長くなればそれ相応のコストになってはきます。しかしながら、リクエスト数が多くないならどう考えてもサーバレスの方がコスト優位と感じる状況です。

cerebriumの使い方

cerebriumは使うのが異常なほど簡単というのも僕が気に入っている点です。
使うのに必要なファイルは主に3つ

cerebrium.toml: 設定ファイル。インスタンスタイプなどの諸条件を記載する。
requirements.txt: おなじみpipのrequirementsを書いておくと、必要なライブラリをインストールしてくれる。
main.py: 推論のコード。ペライチでよい。

まずお手元の環境で

pip install --upgrade cerebrium

して、

cerebrium login あなたの環境

ログインして、

cerebrium init あなたのプロジェクト

でプロジェクトを作成します。

その後、必要なライブラリのrequirements.txtと

main.py

from typing import Optional

from pydantic import BaseModel
from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(
    "stabilityai/japanese-stablelm-instruct-gamma-7b"
)
model = AutoModelForCausalLM.from_pretrained(
    "stabilityai/japanese-stablelm-instruct-gamma-7b",
    torch_dtype="auto",
)
model.eval()
model = model.to("cuda")


########################################
# User-facing API Parameters
########################################
class Item(BaseModel):
    question: str
    temperature: Optional[float] = 1
    top_p: Optional[float] = 0.95
    max_new_tokens: Optional[int] = 256


def build_prompt(user_query, inputs="", sep="\n\n### "):
    sys_msg = "以下は、タスクを説明する指示と、文脈のある入力の組み合わせです。要求を適切に満たす応答を書きなさい。"
    p = sys_msg
    roles = ["指示", "応答"]
    msgs = [": \n" + user_query, ": \n"]
    if inputs:
        roles.insert(1, "入力")
        msgs.insert(1, ": \n" + inputs)
    for role, msg in zip(roles, msgs):
        p += sep + role + msg
    return p


def generate_response(params):
    # Infer with prompt without any additional input
    user_inputs = {
        "user_query": "質問に答えてください。",
        "inputs": params.question,
    }
    prompt = build_prompt(**user_inputs)

    input_ids = tokenizer.encode(prompt, add_special_tokens=True, return_tensors="pt")

    tokens = model.generate(
        input_ids.to(device=model.device),
        max_new_tokens=params.max_new_tokens,
        temperature=params.temperature,
        top_p=params.top_p,
        do_sample=True,
    )

    out = tokenizer.decode(
        tokens[0][input_ids.shape[1]:], skip_special_tokens=True
    ).strip()
    print(out)
    return out


#######################################
# Prediction
#######################################
def predict(item, run_id, logger):
    item = Item(**item)
    result = generate_response(params=item)
    return {"Prediction": result}

こんな感じの推論コードを用意します。
最後に

cerebrium.toml

[cerebrium.build]
predict_data = "{\"question\": \"日本の首都はどこですか\"}"
force_rebuild = false
hide_public_endpoint = false
disable_predict = false
log_level = "INFO"
disable_confirmation = true

[cerebrium.deployment]
name = "your-first-model"
python_version = "3.10"
include = "[./*, main.py, requirements.txt]"
exclude = "[./.*, ./__*]"

[cerebrium.hardware]
gpu = "AMPERE_A5000"
cpu = 2
memory = 64
gpu_count = 1

[cerebrium.scaling]
min_replicas = 0
max_replicas = 1
cooldown = 90

このような設定ファイルを用意したら

cerebrium deploy --name your-first-model --config-file ./cerebriume.toml

これでデプロイが実行できます。
めちゃくちゃ簡単じゃないですか？

設定のポイント

先ほどの設定ファイルではいくつかポイントがあります。
1つはGPUの指定。
ここで使いたいGPUを指定することが可能です。
例では7BだったのでA5000 x 1で指定していますが、70Bの場合はA6000 x 4やA6000 x 2( load_in_8bit )が選択しになってくるでしょう。
利用するGPUの数が増えれば当然料金も増えるので気をつけてください。

2点目はレプリカ数とクールダウンの指定です。
min_replicas を1以上にすると、リクエストがなくてもcold状態にならず常時稼働します。
一気に料金が発生するので気をつけてください。
また、 cooldown　はリクエストがなかった場合にどれくらいの時間でcold状態に移行するかを指定できるものです。
短めにすることでコストを抑えられますが、リクエストの頻度などで調節することをおすすめします。

ここがすごいよcerebrium

僕が気に入っている点はデプロイの容易さだけではありません。
ほかにも「ここがすごいよcerebrium」な点はあります。

まず、モデルロードが早い点です。
普通に考えて70BのモデルをGPUメモリ上にロードするのはすごく時間がかかり、数分かかるのが普通のはずなんですが、cold状態でリクエストしてもわりとすぐにロードされて推論が始まります。
仕組みはいまいちわかってませんが、たとえば自分でEC2を使ってサーバレス風の仕組みをつくろうとしても、停止状態のEC2サーバを起こしてモデルロードするだけで時間がかかってしまい役に立たないと思います。
その点、cerebriumはモデルロードが早いのでいいなと思っています^[2]。

次に、非同期推論できる点です。
cerebriumにデプロイすることでつくられるエンドポイントは、通常だと3分でタイムアウトしてしまいます。
しかし、今回のように70Bのモデルだと推論に3分以上かかることもあり、レスポンスが間に合いません。(coldからだとモデルロードもあるのでさらに時間的に厳しい)
それを解決できるのが　webhook_endpoint オプションです。

これを使うと、まず推論リクエストに対して202とジョブIDのような run_id というものが返されます。
その後、サーバ上で粛々と推論処理が実行され、推論が完了すると webhook_endpoint で指定されたURLに結果がPOST送信されます。
そこには先ほどの run_id と推論の結果が入っているので、 run_id をもとに元々のリクエストと突合して結果を返す処理をつくることができます。
これは本当にすごく便利な機能だと思っていて、サーバレスなうえに非同期処理までサポートされているなんて、こちらではインフラやタイムアウトをなにも気にしなくていいようなものです。
70Bのような推論に長時間かかるものでも、気にせずホスティングして結果はあとからAPI Gatewayなどで受け取ってLambdaでいい感じに整形して元のところに返してやったりすればいいだけなので、非常に楽です^[3]。

実はcerebrium一番のお気に入りポイントはこの非同期処理部分だったりします。

まとめ

そんな感じで、cerebriumを使うといろいろ重たい70Bのモデルもいい感じにホストできますよという話でした。
こんなこと書いて、cerebriumのGPUが枯渇すると困るんですけどね()

というわけで、Spiral.AIでは70Bモデルをホストするくらい機械学習の社会実装と向き合っています！
僕自身、Inf1での成功だけでなく今回のようにさまざまなモデルを安定的に提供できるようにしていけるように引き続き頑張っていきたいと思っています。
もしSpiral.AIに興味がありましたら、お気軽にお声がけください。

脚注

別に実際のリソース状況を正確に聞いたわけではない。使っている感じ、リソース枯渇で詰まったりはしていないので潤沢なのかなと。あと、サポートに聞いたらA6000は枯渇してないよって言われた。なおA100はすでに枯れている。 ↩︎
本当はTensorizerに対応しているのでもっと早くできると思うのですが、僕のセンスが足りずまだ試せていないので気になっています。 ↩︎
LLMはリアルタイム性のあるチャット用途だけではないですからね。非同期でも安定してホスティングして使えると役に立つシーンはあります。 ↩︎

SpiralAIテックブログPublication

実在する芸能人との会話ができる日本初のAIサービス「NaomiAI」やカスタムChatGPTを作れる「Spiralbot」を提供するSpiralAI株式会社のテックブログです。