1800億パラメータのLLMをGoogle Colabで動かす

新しくてパラメータの多いLLMが出る度ごとに、大きいモデルGoogle Colabでは動かせないなと思っていました。

しかし、このPetalsという仕組みを使えば、1800億パラメータを持つtiiuae/falcon-180B-chatが使えます。

詳しい仕組みは省略しますが、BitTorrentスタイルで、分散させて計算させます。その結果、ユーザは使用するLLM全てをダウンロードする必要はありません。

今回のコードは、GitHubからGoogle Colabのリンクがありますので、そちらを利用します。

Google Colabで、実行していきます。Google Colabでは、T4 GPUを使用しました。まずは、Petalsをインストールします。

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

次に、分散モデルをインストールします。今回は、falcon-180B-chatを使います。

import torch
from transformers import AutoTokenizer
from petals import AutoDistributedModelForCausalLM

model_name = "tiiuae/falcon-180B-chat "
# You could also use "meta-llama/Llama-2-70b-chat-hf" or any other supported model from 🤗 Model Hub

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

次に、テキスト生成するために、下記を入力します。A cat in French is と書いて、フランス語で猫はに続く文章を生成してもらいます。

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

所感としては、サンプルを実行してみるとわかりますが、1文字書かれるのに凄く時間がかかります。

この生成スピードなら70億程度のパラメータを動かした方が良いという考えもあるのではないかと思われます。

生成スピードが問題だと思いますが、巨大パラメータを動かせる仕組みは素晴らしいです。

Petalsでは、下記のサイトでhealthyと書かれているLLMが利用可能です。

Petals Health Monitor

利用可能なLLMの分散状態