CyberAgentの日本語LLMを動かしてみる

モデルについて

Hugging Face 🤗 上で公開された CyberAgent さんの日本語版 LLM を動かしてみます。CyberAgent さんの Hugging Face 🤗 ページへのリンクを以下に載せておきます。

リンク先へ飛ぶと、以下の画像のように様々なレポジトリが出てきますが、 open-calm-hoge となっているレポジトリが日本語版 LLM のレポジトリになっています。 hoge の部分はモデルの大きさを表しています。（レポジトリとモデルのサイズの対応表）

使用方法

Python の環境を構築しておきます。PyTorch と Transformers が最低限入っていれば動くと思います。

公式に紹介されている実行コード + α (主に変更しそうな部分にコメント)

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("cyberagent/open-calm-7b", device_map="auto", torch_dtype=torch.float16)
tokenizer = AutoTokenizer.from_pretrained("cyberagent/open-calm-7b")

inputs = tokenizer("AIによって私達の暮らしは、", return_tensors="pt").to(model.device)  # 入力するプロンプト.
with torch.no_grad():
    tokens = model.generate(
        **inputs,
        max_new_tokens=64,  # 生成する長さ. 128 とかでも良い.
        do_sample=True,
        temperature=0.7,  # 生成のランダム性. 高いほど様々な単語が出てくるが関連性は下がる.
        pad_token_id=tokenizer.pad_token_id,
    )
    
output = tokenizer.decode(tokens[0], skip_special_tokens=True)
print(output)

以下のモデルを読み込む際の部分のオプションについて少し紹介します。

モデルを読み込む部分

model = AutoModelForCausalLM.from_pretrained("cyberagent/open-calm-7b", device_map="auto", torch_dtype=torch.float16)

"cyberagent/open-calm-7b" となっている部分を、レポジトリとモデルのサイズの対応表を参照して自分の使いたいモデル名を指定し、変更することができます。（tokenizer の部分も同じものに変更する必要があります。）
また、複数の GPU デバイスを持っている人向けのオプションですが、 device_map="auto" となっている場合は、それぞれの GPU デバイスに均等にモデルがロードされますが、 device_map="sequential" とすることで、1つの GPU デバイスにロードさせることができます。（使用したい GPU デバイスが決まっていれば、os.environ["CUDA_VISIBLE_DEVICES"] = "0"　のように、プログラムの一番最初で指定してあげると良いかもしれません。）
torch_dtype=torch.float16 となっている部分は、モデルを半精度型の浮動小数点数で読み込んでいます。

一番大きい 7B のモデルを torch.float16 型で使用すると、 14~15 GB の GPU メモリを使用しました。
ですが、さらにビット数を落としてモデルを使用することもできます。
Accelerate というライブラリと bitsandbytes というライブラリを追加でインストールすることで、 torch.int8 型で推論させることができます。
これらのライブラリをインストールすれば、モデルを読み込む部分を以下のように変更するだけで、さらに半分の 8bit へと量子化してくれます。

少し変えたモデルを読み込む部分

model = AutoModelForCausalLM.from_pretrained("cyberagent/open-calm-7b", device_map="auto", torch_dtype=torch.int8, load_in_8bit=True)  # torch_dtype=torch.int8 のところが他の型であっても override されます.

こうすることで、同じ 7B のモデルでも、使用する GPU メモリは 8 GB 前後に抑えることができます。（生成の質は若干ですが下がるかもしれません。）

日本語の LLM を手軽に手元で動かせるようになるのはとても嬉しいですね！
CyberAgent さん公開ありがとうございます👏👏👏