背景・目的

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だったので、傾き・切片共に大きめだった。