背景・目的
- OSSのLLMをGPUを使って処理するにあたって、モデルのパラメータ数によって必要なVRAMの量が変わる。
- 実際に使ってみると、入力トークン・出力トークン数によってもVRAM利用量が変わるし、処理時間もトークン数によって違うことがわかってきた。
- 実際に計測してみて、定式化したい
実験内容
- モデルを固定して、入力トークン数・出力トークン数と利用VRAMと生成速度の関係を測定する
- 上記の内容はGPUの単体利用と複数利用で変化するのか？を測定する
前提
| 条件 | 内容 |
|---|---|
| 利用するGPU | NVIDIA GeForce RTX 2080 2台 （VRAM 8GB x 2) |
| NVIDIA Driver | 470.199.02 |
| CUDA | 11.8 |
| モデル | rinna/japanese-gpt-neox-3.6b-instruction-ppo |
| 量子化 | nf4で4bit量子化 |
モデルの読み込み
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
model_name = "rinna/japanese-gpt-neox-3.6b-instruction-ppo"
tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=True)
model = AutoModelForCausalLM.from_pretrained(
model_name,
load_in_4bit=True,
device_map="auto",
max_memory={0: "7GB",1: "7GB" },
quantization_config=BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
),
)
実験1: 入出力トークン数とVRAMと生成速度の関係を測定する
実験内容
下記のパターンでVRAM利用量と生成速度を計測する
| 項目 | パターン |
|---|---|
| 入力トークン数 | 100, 500, 1000 |
| 最大出力トークン数 | 64, 128, 256 |
GPUは複数利用する
VRAMの計測方法
こちらの記事を参考にnvidia-smiの値を取得する
import subprocess
DEFAULT_ATTRIBUTES = (
"index",
"uuid",
"name",
"timestamp",
"memory.total",
"memory.free",
"memory.used",
"utilization.gpu",
"utilization.memory",
)
def get_gpu_info(nvidia_smi_path="nvidia-smi", keys=DEFAULT_ATTRIBUTES, no_units=True):
nu_opt = "" if not no_units else ",nounits"
cmd = "%s --query-gpu=%s --format=csv,noheader%s" % (
nvidia_smi_path,
",".join(keys),
nu_opt,
)
output = subprocess.check_output(cmd, shell=True)
lines = output.decode().split("\n")
lines = [line.strip() for line in lines if line.strip() != ""]
return [{k: v for k, v in zip(keys, line.split(", "))} for line in lines]
時間の計測方法
下記のように、テキストを生成するgenerate関数部分だけにかかった時間を計測する。
入力トークンの計算時間・出力トークンのでコード処理は、計測時間に含めない。（いずれも300ms程度であり、generate関数の部分と比較して小さかったため）
with open(sys.argv[1]) as f:
input_text = f.read()
tokens = create_input_tokens(input_text)
results = []
for input_token_num in [100, 500, 1000]:
for max_output_tokens in [64, 128, 256]:
start = time()
out = generate(
tokens[:, :input_token_num], max_new_tokens=max_output_tokens
)
t = time() - start
info = get_gpu_info()
mem = sum([int(u["memory.used"]) for u in info])
results.append(
(input_token_num, len(out[0]), len(out[0]) - input_token_num, t, mem)
)
実験結果
| 入力トークン数 | 出力トークン数 | 入出力トークン数 | 生成速度(s) | VRAM |
|---|---|---|---|---|
| 100 | 64 | 164 | 9.96 | 4376 |
| 100 | 128 | 228 | 18.56 | 4502 |
| 100 | 256 | 356 | 37.02 | 4688 |
| 500 | 64 | 564 | 9.92 | 5060 |
| 500 | 128 | 628 | 19.18 | 5060 |
| 500 | 102 | 602 | 15.42 | 5060 |
| 1000 | 64 | 1064 | 10.79 | 6078 |
| 1000 | 128 | 1128 | 20.12 | 6078 |
| 1000 | 256 | 1256 | 38.87 | 6078 |
1. 入力トークン数と処理時間・VRAMの関係について
- 入力トークン数が増えるほど利用するVRAMが増えていることがわかる
- しかし、処理時間については入力トークン数に依存していなさそうに見える
2. 出力トークン数と処理時間・VRAMの関係について
- VRAMについては、出力トークン数以外の要因によって決まっているように見える
- しかし、処理時間については、出力トークン数に依存して増えていそうに見える
3. 入出力トークン数と処理時間・VRAMの関係について
- 入出力トークン数（=入力トークン数+出力トークン数)が増えるほど、VRAMの利用量が増えているように見える。先ほどの、出力トークン数とVRAMとの関係については、入力トークン数の数の方が影響が大きかったため見えづらかったが、入出力トークン数が小さいときは出力トークン数もVRAMに影響していそうである。
- とはいえ、トークン数合計が600あたりの時には、出力トークン数の影響がVRAM利用量に影響していなさそうに見える（多めにVRAMを確保していて、その中で賄えているのかもしれない）
- 処理時間については、やはり出力トークン数の影響が大きいように見える
近似
- 入出力トークン数とVRAMの関係を直線で近似した結果
y=1.7085x+4082.8658- 係数 1.7 (MB/トークン数)で増えていく
結論
- VRAMは入出力のトークン数に依存して増える。
- 処理時間は出力トークン数に依存して増える。
実験2 GPUを単体利用と複数利用で処理時間が変わるか？
accelerateのライブラリを使って複数のGPUでモデルを動かしているが、一つのGPUで動かした方が通信オーバーヘッドが小さくなるので生成速度が早くなるのでは？と思い、調べてみました
実験内容
下記のパターンでVRAM利用量と生成速度を計測する
| 項目 | パターン |
|---|---|
| 入力トークン数 | 100, 500, 1000 |
| 最大出力トークン数 | 64, 128, 256 |
GPUは一つだけ利用する。前回の実験で、上記のパターンであれば、GPUVRAMは7GBに収まるので1台で処理できるはず。
モデルの読み込み部分を少し修正し、一つだけGPUを利用するようにする。
model = AutoModelForCausalLM.from_pretrained(
checkpoint,
load_in_4bit=True,
device_map="auto",
# max_memory={0: "7GB",1: "7GB" },
max_memory={0: "7GB"},
quantization_config=BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
),
)
実験結果
| 入力トークン数 | 出力トークン数 | 入出力トークン数 | 生成速度(s) | VRAM |
|---|---|---|---|---|
| 100 | 64 | 164 | 9.54 | 3393 |
| 100 | 128 | 228 | 18.43 | 3521 |
| 100 | 174 | 274 | 25.05 | 3583 |
| 500 | 64 | 564 | 9.93 | 4079 |
| 500 | 128 | 628 | 19.34 | 4079 |
| 500 | 256 | 756 | 38.16 | 4203 |
| 1000 | 64 | 1064 | 10.98 | 4997 |
| 1000 | 128 | 1128 | 20.52 | 4997 |
| 1000 | 256 | 1256 | 39.63 | 4997 |
単体利用の時と複数利用の時、それぞれについて、出力トークン数と処理時間の関係を見てみたところ、今回の条件下だと処理時間にほとんど差がないことがわかった。
細かいところはわからないが、ネックが通信とは別のところにあると思われる。
一時関数で近似すると上記のようになりました。
y=0.1446x+0.8596だったので、出力トークン数1つにつき、処理時間が0.15秒増える計算になりそうです
まとめ
今回は、トークン数とVRAM利用量・処理時間の関係について確認してみました。
- VRAM利用量はモデルのパラメータ数だけではなく、入出力トークン数にも依存して増える
- 処理時間は、出力トークン数に依存して増える
- 処理時間は、GPUの数を減らしても変わらない
ということがわかりました。
今回、トークン数を増やすためには沢山のVRAMだろうという仮説が確認できたことは良かったのと、
処理時間は入力トークン数に依存しないということは意外でした。
実験結果について、毎回GPUのメモリをリセットしていないので、その影響が出てしまっている可能性もありそう。（トークン数が少ない順に処理しているので、確保したVRAMが解放されない場合にトークン数増加に伴いVRAMが単調増加するのは当たり前なので・・）少し時間はかかるが、毎回モデルを読み込み直すアプローチで、また計測し直してみようと思います。
[追記] -> 試したところ、あまり結果は変わりませんでした。
また、他のモデルでも計測してみたいと思います。
追記
star coder (15.5B) で試した結果
-
入出力トークンとVRAMの関係は、だいたい同じだったが一次関数で近似は難しそうだった
-
トークン数1000を超えたあたりで、一段階必要なVRAM量が増えるようなグラフになった
-
出力トークン数を64に固定し、入力トークン数を50刻みで増やして計測した結果はこちら
線形に増加しているというよりは、段々に増えているように見える
-
-
出力トークンと生成時間は線形に増加していそう
-
y = 0.47x + 9.08だったので、傾き・切片共に大きめだった。
-