株式会社サイバーエージェント が開発した大規模言語モデル open-calm-7b の動作確認を Windows かつ GeForce RTX 3060 を搭載したローカル PC で行ってみました。
確認には以下の記事を参考にさせていただきました。
大規模言語モデルに興味のある方は是非是非このお二方のいいね!とフォローをお願いします。
- npaka 氏: Google Colab で PEFT による大規模言語モデルのファインチューニングを試す
- m__k 氏: OpenCALM-7B を LoRA で instruction tuning するための実装解説 / QLoRA の実装も紹介
open-calm-7b の動かし方については m__k 氏が 既にまとめておられます ので、この記事の価値はそれを GeForce RTX 3060 で試してみたらこうなった、程度のものしかありません。ご了承ください。
世の中で最も使われているであろう大規模言語モデルとして GPT が挙げられますが、組織のセキュリティポリシーによっては機密情報を入力できないという問題があります。そこで、組織内で立てたサーバに大規模言語モデルを入れて GPT を使ったワークフローと同等のことができないか?と考え、株式会社サイバーエージェントが開発した日本語を生成できるモデルである open-calm-7b を試しに動かしてみる事にしました。さらに、なるべく運用コストを下げたかったため、なるべくローカル PC で推論からファインチューニングまで行えないか調べる事にしました。
また、open-calm-7b の性能について所感を先に述べておくと、GPT ほどの汎用性と正確性は無く、用途を選んで使用するか、ファインチューニングを行って特化型のモデルにしていくかのいずれかが必要になりそうです。(GPTの方がモデルのパラメータ数が非常に多いので当然の結果ではあるのですが。)
文脈に基づく判定や推論のような正確さが求められる用途は厳しい反面、要約やタイトル生成などは何度もリトライして良いものを選べばいいため比較的使いやすいのではないかと思います。
動作環境
今回 open-calm-7b を動かす環境は Dell の GeForce RTX 3060 を搭載したゲーミング PC です。
構成は以下の通りです。
- Dell Alienware Aurora R12 (ミドルタワー)
- Intel Core i7-11700F
- 32 GB RAM
- 2 TB HDD
- Windows 11
- Python 3.10.6
- Git Bash
- Visual Studio Code
- Jupyter 拡張
- GeForce RTX 3060 (12 GB VRAM)
- NVIDIA CUDA Runtime 11.7
この構成の中で Windows OS と 12 GB VRAM の 2 点については、モデルを普通に動かす上で大きなハードルになります。
もし諸々のトラブルを避けたい場合はコストをかけてでも、フルタワー PC、Linux、GeForce RTX 4090 1 等のローカル PC 構成か、Google Colab や Paperspace の高スペックインスタンスを使うことをオススメします。
事前準備
依存パッケージのインストール
動作確認に必要なパッケージをインストールします。
まずは PyTorch をインストールします。
インストール完了後 CUDA が有効になっていることを確認してください。
pip install -U torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu117
import torch
print(f"torch.cuda.is_available(): {torch.cuda.is_available()}")
print(f"torch.version.cuda: {torch.version.cuda}")
torch.cuda.is_available(): True
torch.version.cuda: 11.7
次に PyTorch 以外のパッケージをインストールします。
transformers
や peft
については m__k 氏のやり方 に倣って QLoRA でファインチューニングが行えるようにパッケージの参照先を設定しています。
QLoRA でファインチューニングを行うには、重みパラメータを量子化するために bitsandbytes
が必要ですが、公式リポジトリでは Windows 向けのバイナリを提供していないため、Windows 向けのバイナリを提供してくださっているリポジトリからインストールしています。
なお、QLoRA によるチューニングを行う理由は、RTX 3060 のようにモデルサイズに対して少ない VRAM でも読み込めるようにするためです。bitsandbytes
による 8bit, 4bit 量子化が無ければ読み込むことすらままなりません。float32 でモデルを読み込むと単純計算で 28 GB の VRAM を消費することになるため 12 GB の VRAM では容量オーバーとなります。これが 4bit 量子化を行うことで、推論時は約 6GB、ファインチューニング時は 11.5GB の VRAM 消費で済みました。
量子化の研究の進みによってコストパフォーマンスが2倍、4倍と段階的に良くなっていますので、より実用的なモデルを使いたい場合は必須のテクノロジかと思います。
量子化技術については以下のページを参考にしています。
- A Gentle Introduction to 8-bit Matrix Multiplication for transformers at scale using Hugging Face Transformers, Accelerate and bitsandbytes
- Making LLMs even more accessible with bitsandbytes, 4-bit quantization and QLoRA
pip install -U https://github.com/jllllll/bitsandbytes-windows-webui/raw/main/bitsandbytes-0.39.0-py3-none-any.whl
pip install -U git+https://github.com/huggingface/transformers.git
pip install -U git+https://github.com/huggingface/peft.git
pip install -U git+https://github.com/huggingface/accelerate.git
pip install -U datasets ipywidgets sentencepiece
念のためインストールされているパッケージの内容を確認します。
特に標準パッケージリポジトリ以外からインストールしているものについては、正しくインストールされているか注意が必要です。
pip freeze
open-calm-7b による生成テスト
HuggingFace から open-calm-7b を読み込んでテキスト生成をテストします。
import torch
import transformers
まずは open-calm-7b モデルを読み込みます。初めて読み込む際は HuggingFace からのファイルダウンロードが発生するため時間がかかります。モデルのファイルサイズは約12GBになるためディスクドライブの容量に注意してください。
なお、この段階では 8bit 量子化モードで読み込んでいます。from_pretrained
の load_in_8bit
のパラメータに True
を指定することで bitsandbytes
による 8bit 量子化が行われ、VRAM に展開されるパラメータのサイズが圧縮されます。
base_model_name = "cyberagent/open-calm-7b"
tokenizer = transformers.AutoTokenizer.from_pretrained(base_model_name)
base_model = transformers.AutoModelForCausalLM.from_pretrained(base_model_name, device_map='auto', load_in_8bit=True)
print(base_model)
GPTNeoXForCausalLM(
(gpt_neox): GPTNeoXModel(
(embed_in): Embedding(52224, 4096)
(layers): ModuleList(
(0-31): 32 x GPTNeoXLayer(
(input_layernorm): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)
(post_attention_layernorm): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)
(attention): GPTNeoXAttention(
(rotary_emb): RotaryEmbedding()
(query_key_value): Linear8bitLt(in_features=4096, out_features=12288, bias=True)
(dense): Linear8bitLt(in_features=4096, out_features=4096, bias=True)
)
(mlp): GPTNeoXMLP(
(dense_h_to_4h): Linear8bitLt(in_features=4096, out_features=16384, bias=True)
(dense_4h_to_h): Linear8bitLt(in_features=16384, out_features=4096, bias=True)
(act): GELUActivation()
)
)
)
(final_layer_norm): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)
)
(embed_out): Linear(in_features=4096, out_features=52224, bias=False)
)
テキスト生成を実行します。プロンプトには『吾輩は猫である』の冒頭の文章を指定します。
高い確率で『吾輩は猫である』の続きの文章が生成されます。
prompt = """吾輩は猫である。名前はまだ無い。どこで"""
inputs = tokenizer(prompt, return_tensors="pt").to(base_model.device)
with torch.no_grad():
tokens = base_model.generate(
**inputs,
max_new_tokens=64,
do_sample=True,
temperature=0.7,
top_p=0.9,
repetition_penalty=1.05,
pad_token_id=tokenizer.pad_token_id,
)
outputs = tokenizer.batch_decode(tokens)
output = outputs[0]
print(output)
nvidia-smi
コマンドを使って VRAM や温度の状況を確認することができます。
モデルを読み込んだ状態では約 8 GB の VRAM を消費しました。
!nvidia-smi
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.98 Driver Version: 535.98 CUDA Version: 12.2 |
|-----------------------------------------+----------------------+----------------------+
| GPU Name TCC/WDDM | Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|=========================================+======================+======================|
| 0 NVIDIA GeForce RTX 3060 WDDM | 00000000:02:00.0 On | N/A |
| 0% 41C P8 16W / 170W | 8502MiB / 12288MiB | 0% Default |
| | | N/A |
+-----------------------------------------+----------------------+----------------------+
QLoRA によるファインチューニング
open-calm-7b は GPT のような超巨大モデルではないため、なんとなく使うだけでは実用に耐えないテキストばかり生成されるかもしれません。
そのため、モデルの利用用途をある程度決めておいて、その用途に合わせてファインチューニングを行うことで実用にも耐えうるモデルになることを期待してファインチューニングを試してみます。
テキスト生成の時と同様に RTX 3060 を使用します。
ファインチューニングの手法として、個人的には画像生成のほうでもなじみのある LoRA を使用します。
さらに 4bit 量子化されたモデルに大してチューニングが可能な QLoRA を用いて、少ない GPU メモリで open-calm-7b がチューニングできるか確認してみます。
import datasets
import peft
import torch
import torch.nn
import transformers
まずはベースとなる open-calm-7b
を読み込みます。その際 4bit 量子化モードを有効化します。
base_model_name = "cyberagent/open-calm-7b"
tokenizer = transformers.AutoTokenizer.from_pretrained(base_model_name)
bnb_config = transformers.BitsAndBytesConfig(
load_in_4bit=True, # 4bit 量子化を有効化します。
bnb_4bit_use_double_quant=True, # ネストされた量子化 (Nested quantization) を有効化します。
bnb_4bit_quant_type="nf4", # 量子化データタイプを設定します。nf4 は 4-bit NormalFloat Quantization のデータ型です。
bnb_4bit_compute_dtype=torch.bfloat16 # 量子化計算時のデータタイプを設定します。
)
base_model = transformers.AutoModelForCausalLM.from_pretrained(base_model_name, device_map='auto', quantization_config=bnb_config)
print(base_model)
GPTNeoXForCausalLM(
(gpt_neox): GPTNeoXModel(
(embed_in): Embedding(52224, 4096)
(layers): ModuleList(
(0-31): 32 x GPTNeoXLayer(
(input_layernorm): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)
(post_attention_layernorm): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)
(attention): GPTNeoXAttention(
(rotary_emb): RotaryEmbedding()
(query_key_value): Linear4bit(in_features=4096, out_features=12288, bias=True)
(dense): Linear4bit(in_features=4096, out_features=4096, bias=True)
)
(mlp): GPTNeoXMLP(
(dense_h_to_4h): Linear4bit(in_features=4096, out_features=16384, bias=True)
(dense_4h_to_h): Linear4bit(in_features=16384, out_features=4096, bias=True)
(act): GELUActivation()
)
)
)
(final_layer_norm): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)
)
(embed_out): Linear(in_features=4096, out_features=52224, bias=False)
)
読み込んだモデルに対して訓練用の加工を行います。
for param in base_model.parameters():
# ベースモデルのパラメータは勾配計算の対象外とします。
param.requires_grad = False
# 学習に使用するレイヤーのパラメータの型を float32 にして精度を上げます。
if param.ndim == 1:
param.data = param.data.to(torch.float32)
# メモリを節約するために、Gradient Checkpointing アルゴリズムを有効化します。
base_model.gradient_checkpointing_enable()
# モデルの重みを固定したままアダプターの重みを調整するために、入力埋め込みグラデーションを有効化します。
base_model.enable_input_require_grads()
class CastOutputToFloat(torch.nn.Sequential):
def forward(self, x):
return super().forward(x).to(torch.float32)
base_model.embed_out = CastOutputToFloat(base_model.embed_out)
print(base_model)
PEFT を使って QLoRA の準備を行います。
PEFT を使えば、少量あるいは追加のパラメータのみをチューニングするいくつかのアルゴリズムを簡単に導入することができます。
peft_config = peft.LoraConfig(
r=8,
lora_alpha=32,
target_modules=["query_key_value"],
lora_dropout=0.05,
bias="none",
fan_in_fan_out=False,
task_type=peft.TaskType.CAUSAL_LM
)
peft_model = peft.get_peft_model(base_model, peft_config)
peft_model.print_trainable_parameters()
トレーニング用のデータセットを用意します。
kunishou/databricks-dolly-15k-ja
のデータセットを元に、最終的な入力を想定したプロンプトを組み立てて訓練用データとしています。
プロンプトフォーマットは m__k氏の記事 をほとんど真似ていますが、日本語のモデルに英語の命令文を入れるとどうなるのか少し興味があったので英語にしています。(英語にしたことによる変化はまだ検証できていません。)
dataset_name = "kunishou/databricks-dolly-15k-ja"
dataset = datasets.load_dataset(dataset_name)
prompt_with_context_format = """The following text is the task instruction and the context for it.
Write a response that satisfies the instruction based on context.
### Instruction:
{instruction}
### Context:
{context}
### Response:
{response}
"""
prompt_no_context_format = """The following text is the task instruction.
Write a response that satisfies the instruction based on context.
### Instruction:
{instruction}
### Response:
{response}
"""
def tokenize(samples):
prompts = []
# データセットの instruction 列と input 列と output 列を組み合わせてプロンプトを組み立てます。
for instruction, input, output in zip(samples["instruction"], samples["input"], samples["output"]):
if input:
prompt = prompt_with_context_format.format(instruction=instruction, context=input, response=output)
else:
prompt = prompt_no_context_format.format(instruction=instruction, response=output)
prompts.append(prompt + tokenizer.eos_token)
result = tokenizer(prompts, padding=False, truncation=True, max_length=2048)
return result
dataset = dataset.map(lambda samples: tokenize(samples), batched=True)
transformers.Trainer
を使って訓練を開始します。
200 ステップの学習に GeForce RTX 3060 で約 1 時間かかりました。
Out Of Memory を極力起こしたくないため、バッチサイズ per_device_train_batch_size
は少なめにしていますが、学習中の VRAM 消費量は 11.5 GB 程度に達しており上限ギリギリの状態でした。
training_args = transformers.TrainingArguments(
per_device_train_batch_size=2,
gradient_accumulation_steps=8,
warmup_steps=100,
max_steps=200,
learning_rate=2e-4,
fp16=True,
num_train_epochs=1,
save_strategy="steps",
save_steps=10,
save_total_limit=10,
output_dir=".checkpoints",
evaluation_strategy="no",
logging_dir="logs",
logging_steps=10,
push_to_hub=False,
)
data_collator = transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False)
trainer = transformers.Trainer(
model=peft_model,
train_dataset=dataset["train"],
args=training_args,
data_collator=data_collator,
)
peft_model.config.use_cache = False
# 学習を途中から再開する場合はここへチェックポイント名を記入します。
checkpoint = None
# checkpoint = "checkpoint-100"
trainer.train(checkpoint)
peft_model.config.use_cache = True
QLoRA のパラメータファイルを .lora
ディレクトリに保存します。
peft_model.save_pretrained(".lora")
QLoRA によって得られたファインチューニング付きで推論テスト
QLoRA によって得られたパラメータを追加で読み込んで、やや難しめの推論テストを行います。
import peft
import torch
import torch.nn
import transformers
open-calm-7b
と QLoRA で得られたパラメータを読み込みます。
base_model_name = "cyberagent/open-calm-7b"
tokenizer = transformers.AutoTokenizer.from_pretrained(base_model_name)
bnb_config = transformers.BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16
)
base_model = transformers.AutoModelForCausalLM.from_pretrained(base_model_name, device_map='auto', quantization_config=bnb_config)
model = peft.PeftModel.from_pretrained(base_model, '.lora')
print(model)
以下はモデルのフォームです。
アテンションレイヤーに LoRA レイヤーが追加されていることを確認します。
PeftModelForCausalLM(
(base_model): LoraModel(
(model): GPTNeoXForCausalLM(
(gpt_neox): GPTNeoXModel(
(embed_in): Embedding(52224, 4096)
(layers): ModuleList(
(0-31): 32 x GPTNeoXLayer(
(input_layernorm): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)
(post_attention_layernorm): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)
(attention): GPTNeoXAttention(
(rotary_emb): RotaryEmbedding()
(query_key_value): Linear4bit(
in_features=4096, out_features=12288, bias=True
(lora_dropout): ModuleDict(
(default): Dropout(p=0.05, inplace=False)
)
(lora_A): ModuleDict(
(default): Linear(in_features=4096, out_features=8, bias=False)
)
(lora_B): ModuleDict(
(default): Linear(in_features=8, out_features=12288, bias=False)
)
(lora_embedding_A): ParameterDict()
(lora_embedding_B): ParameterDict()
)
(dense): Linear4bit(in_features=4096, out_features=4096, bias=True)
)
(mlp): GPTNeoXMLP(
(dense_h_to_4h): Linear4bit(in_features=4096, out_features=16384, bias=True)
(dense_4h_to_h): Linear4bit(in_features=16384, out_features=4096, bias=True)
(act): GELUActivation()
)
)
)
(final_layer_norm): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)
)
(embed_out): Linear(in_features=4096, out_features=52224, bias=False)
)
)
)
読み込んだモデルを使って推論を実行します。
生成時間は出力結果にもよりますが約 20 秒かかります。
家族構成のコンテキストをベースに指定された条件に合う人を列挙するというものですが、結果としては間違った回答がほとんど(というか正解を得られたことが無い)です。
prompt = """The following text is the task instruction and the context for it.
Write a response that satisfies the instruction based on context.
### Instruction:
サザエさん一家の女性を列挙してください。
### Context:
- サザエさん一家はサザエ、フネ、ナミヘイ、マスオ、カツオ、タラオ、ワカメ、タマの7名と1匹で構成されている。
- サザエはフネの娘である。
- フネはナミヘイの妻である。
- マスオはタラオの父である。
- カツオはナミヘイの息子である。
- ワカメはカツオの妹である。
- タマはサザエさん一家の飼い猫である。
### Response:
"""
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
with torch.no_grad():
tokens = model.generate(
**inputs,
max_new_tokens=64,
do_sample=True,
temperature=0.7,
top_p=0.9,
repetition_penalty=1.05,
pad_token_id=tokenizer.pad_token_id,
)
outputs = tokenizer.batch_decode(tokens)
output = outputs[0]
print(output)
女性たちの名前は以下になります:
サザエ・ノリエガ(Faine Suzuenege)
ハナコ・ムラカミ(Hanako Murakami)
イクラ・タマ子(Ikra Tamako)
カオリン・トメ子(Kazun
ちなみに GPT の Playground から同様の質問をしてみると概ね合ってはいるのですが、ナミヘイが女性として列挙されてしまうなどの誤回答が発生します。
Gradio でテスト用 UI を起動する
最後に Gradio を使い、Web ブラウザから open-calm-7b および QLoRA で得られたパラメータを読み込んだモデルの推論テストを行えるようにします。
まずは Gradio パッケージをインストールします。
pip install -U gradio
import gradio as gr
import peft
import torch
import torch.nn
import transformers
モデルと QLoRA のパラメータを読み込みます。
base_model_name = "cyberagent/open-calm-7b"
tokenizer = transformers.AutoTokenizer.from_pretrained(base_model_name)
bnb_config = transformers.BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16
)
base_model = transformers.AutoModelForCausalLM.from_pretrained(base_model_name, device_map='auto', quantization_config=bnb_config)
# QLoRA により学習した追加のパラメータがある場合はこちらを実行します。
model = peft.PeftModel.from_pretrained(base_model, '.lora')
# QLoRA により学習した追加のパラメータが無い場合はこちらを実行します。
# model = base_model
print(model)
Gradio を開始します。
Web ブラウザでターミナルに出力された URL へアクセスすることで、テスト用 Web ページを表示できます。
左側のペインに表示されたテキスト入力部分に補完したいプロンプトを入力して送信することで、右側のペインに補完結果が表示されます。
def complement(prompt, temperature):
inputs = tokenizer(prompt, return_tensors='pt').to(model.device)
with torch.no_grad():
tokens = model.generate(
**inputs,
max_new_tokens=128,
do_sample=True,
temperature=temperature,
top_p=0.9,
repetition_penalty=1.05,
pad_token_id=tokenizer.pad_token_id,
)
outputs = tokenizer.batch_decode(tokens)
output = outputs[0]
return output
demo = gr.Interface(
complement,
[
gr.Textbox(lines=10, label="Prompt"),
gr.Number(value=0.7, label='Temperature'),
],
gr.Textbox(label="Output"),
)
demo.launch()
まとめ
とりあえず、ローカル PC で open-calm-7b を動かすという当初の目標は達成しました。
正直 GeForce RTX 3060 で動かせるのはかなり大きい収穫でした。
これでローカルで気兼ねなく学習とテストを回すことができます。
より詳細な open-calm-7b の性能検証や、ファインチューニングの効果検証はこれから行う予定です。
ただ、やはり GPT の性能が圧倒的だというのを再確認して、果たして 7b サイズの日本語モデルで実用的な使い方はできるんだろうか?と若干不安になってきた次第です。
-
もちろん A100 のように、より高価な GPU でも OK です。 ↩