見出し画像

Rinna 3.6B の量子化とメモリ消費量

npaka

「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 bitsandbytes

2-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 s

2-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 s

2-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 s

2-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

関連



この記事が気に入ったらサポートをしてみませんか？