Rinna 3.6B の量子化とメモリ消費量
「Google Colabでの「Rinna 3.6B」の量子化とメモリ消費量を調べてみました。
1. 量子化とメモリ消費量
「量子化」は、LLMのメモリ消費量を削減するための手法の1つです。通常、メモリ使用量が削減のトレードオフとして、LLMの精度が低下します。
AutoTokenizer.from_pretrained()の以下のパラメータを調整します。
・torch_dtype : PyTorchのデータ型の指定。(torch.float16など)
・load_in_8bit : 8-bit精度でモデル読み込み。(bitsandbytes)
・load_in_4bit : 8-bit精度でモデル読み込み。(bitsandbytes)
・device_map : デバイス指定。(accelerate)
2. Colabでの確認
Colabでの実行手順は、以下のRLHFモデルと同様です。
速度は、モデル読み込み後、2回目の推論で速度計測しています。
(1回目は時間かかる)
load_in_8bit / load_in_4bit / device_map を利用するために、以下のパッケージもインストールする必要があります。
!pip install accelerate bitsandbytes2-1. float32
設定なしで動作確認します。
◎ 設定
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
# トークナイザーとモデルの準備
tokenizer = AutoTokenizer.from_pretrained(
"rinna/japanese-gpt-neox-3.6b-instruction-ppo",
use_fast=False
)
model = AutoModelForCausalLM.from_pretrained(
"rinna/japanese-gpt-neox-3.6b-instruction-ppo",
device_map="auto",
)◎ 結果
私は、ほむらが最も可愛いと思います。彼女はとても優しいです。彼女はまどかに魔法を教え、友達でした。彼女はとても強い力を持っています。彼女の外見もとても美しいです。彼女はとてもかわいくて強いです。彼女はとても魅力的だと思います。彼女はとても素晴らしいキャラクターです。彼女はとても美しく、優しくて強いです。彼女はとても良い親友です。彼女はとても強い存在感があり、とても賢いです。彼女はとても素晴らしく、強い力を持っています。彼女はとても素直で優しいです。彼女はとても強い存在感があり、とてもやさしいです。彼女はとても優しい心を持っているので、とても優れた
CPU times: user 12.6 s, sys: 145 ms, total: 12.7 s
Wall time: 4.63 s2-2. float16
「torch_dtype=torch.float16」で動作確認します。
◎ 設定
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
# トークナイザーとモデルの準備
tokenizer = AutoTokenizer.from_pretrained(
"rinna/japanese-gpt-neox-3.6b-instruction-ppo",
use_fast=False
)
model = AutoModelForCausalLM.from_pretrained(
"rinna/japanese-gpt-neox-3.6b-instruction-ppo",
torch_dtype=torch.float16,
device_map="auto",
)◎ 結果
ほむらです。彼女はとても魅力的で、とても複雑なキャラクターとして描かれています。彼女の主な焦点は、他の魔法少女たちとの友情や対立にあります。彼女は孤独な戦いを続けながら、多くの友人や仲間を持ち、強い意志を持っています。彼女はとても誇り高く、強い存在感があり、とてもかわいくて魅力的です。</s>
CPU times: user 6.88 s, sys: 91.5 ms, total: 6.97 s
Wall time: 5.27 s2-3. 8-bit量子化
「load_in_8bit=True」で動作確認します。
◎ 設定
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
# トークナイザーとモデルの準備
tokenizer = AutoTokenizer.from_pretrained(
"rinna/japanese-gpt-neox-3.6b-instruction-ppo",
use_fast=False
)
model = AutoModelForCausalLM.from_pretrained(
"rinna/japanese-gpt-neox-3.6b-instruction-ppo",
load_in_8bit=True,
device_map="auto",
)◎ 結果
ほむらです。彼女はとてもかわいく、才能があり、強い意志を持っているので、とても素敵です。彼女は友情や優しさ、そして悲しみも持っているので、素晴らしいキャラクターです。彼女はとても魅力的で、強い存在感があります。私は彼女をとても気に入っています。</s>
CPU times: user 12.8 s, sys: 0 ns, total: 12.8 s
Wall time: 12.9 s2-4. 4-bit量子化
「load_in_4bit=True」で動作確認します。
◎ 設定
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
# トークナイザーとモデルの準備
tokenizer = AutoTokenizer.from_pretrained(
"rinna/japanese-gpt-neox-3.6b-instruction-ppo",
use_fast=False
)
model = AutoModelForCausalLM.from_pretrained(
"rinna/japanese-gpt-neox-3.6b-instruction-ppo",
load_in_4bit=True,
device_map="auto",
)◎ 結果
私はまどか☆マギカで、ほむらが最も可愛いと思います。彼女は非常に魅力的で優しく、自尊心の強い女の子です。彼女はとても強い個性を持っています。彼女の主な役割は、人々に自分の世界への理解と共感を与えることです。彼女は、街のすべての物事に意味があるという考えを広めています。また、彼女は重要な存在であり、物語の中心的なキャラクターとしての役割を果たしています。</s>
CPU times: user 13.9 s, sys: 9.09 ms, total: 13.9 s
Wall time: 14.5 s