LLMのモデルの生成速度とVRAM利用量について調査


# 背景・目的

+ OSSのLLMをGPUを使って処理するにあたって、モデルのパラメータ数によって必要なVRAMの量が変わる。
+ 実際に使ってみると、入力トークン・出力トークン数によってもVRAM利用量が変わるし、処理時間もトークン数によって違うことがわかってきた。
+ 実際に計測してみて、定式化したい

# 実験内容

1. モデルを固定して、入力トークン数・出力トークン数と利用VRAMと生成速度の関係を測定する
1. 上記の内容は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量子化 |

モデルの読み込み

```python
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の計測方法

[こちらの記事](https://qiita.com/tomotaka_ito/items/1da001c98b46ecf28ec7)を参考にnvidia-smiの値を取得する

```python
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関数の部分と比較して小さかったため）

```python
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の関係について

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/144584/550e9e33-7a3d-3405-c807-a87b045f2c58.png)

+ 入力トークン数が増えるほど利用するVRAMが増えていることがわかる
+ しかし、処理時間については入力トークン数に依存していなさそうに見える

##### 2. 出力トークン数と処理時間・VRAMの関係について

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/144584/c71a40b8-9472-c0cb-d228-0b9903e8f42d.png)

+ VRAMについては、出力トークン数以外の要因によって決まっているように見える
+ しかし、処理時間については、出力トークン数に依存して増えていそうに見える

##### 3. 入出力トークン数と処理時間・VRAMの関係について

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/144584/6b402c00-122e-9969-e82c-7b09eaa0ad7b.png)


+ 入出力トークン数（=入力トークン数+出力トークン数)が増えるほど、VRAMの利用量が増えているように見える。先ほどの、出力トークン数とVRAMとの関係については、入力トークン数の数の方が影響が大きかったため見えづらかったが、入出力トークン数が小さいときは出力トークン数もVRAMに影響していそうである。
+ とはいえ、トークン数合計が600あたりの時には、出力トークン数の影響がVRAM利用量に影響していなさそうに見える（多めにVRAMを確保していて、その中で賄えているのかもしれない）
+ 処理時間については、やはり出力トークン数の影響が大きいように見える

##### 近似

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/144584/0325c73a-decb-ad63-47d8-1eb96fb23a79.png)

+ 入出力トークン数と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を利用するようにする。

```python
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|



![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/144584/6110bf67-4304-8d72-9fc6-11019955abbc.png)


単体利用の時と複数利用の時、それぞれについて、出力トークン数と処理時間の関係を見てみたところ、今回の条件下だと処理時間にほとんど差がないことがわかった。
細かいところはわからないが、ネックが通信とは別のところにあると思われる。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/144584/e12612b0-e10d-70e3-142b-9d308e9f04e7.png)

一時関数で近似すると上記のようになりました。

`y=0.1446x+0.8596`だったので、出力トークン数1つにつき、処理時間が0.15秒増える計算になりそうです


# まとめ

今回は、トークン数とVRAM利用量・処理時間の関係について確認してみました。

+ VRAM利用量はモデルのパラメータ数だけではなく、入出力トークン数にも依存して増える
+ 処理時間は、出力トークン数に依存して増える
+ 処理時間は、GPUの数を減らしても変わらない

ということがわかりました。
今回、トークン数を増やすためには沢山のVRAMだろうという仮説が確認できたことは良かったのと、
処理時間は入力トークン数に依存しないということは意外でした。


実験結果について、毎回GPUのメモリをリセットしていないので、その影響が出てしまっている可能性もありそう。（トークン数が少ない順に処理しているので、確保したVRAMが解放されない場合にトークン数増加に伴いVRAMが単調増加するのは当たり前なので・・）少し時間はかかるが、毎回モデルを読み込み直すアプローチで、また計測し直してみようと思います。
[追記] -> 試したところ、あまり結果は変わりませんでした。


また、他のモデルでも計測してみたいと思います。

### 追記


#### [star coder (15.5B)](https://huggingface.co/bigcode/starcoder) で試した結果

+ 入出力トークンとVRAMの関係は、だいたい同じだったが一次関数で近似は難しそうだった
    + トークン数1000を超えたあたりで、一段階必要なVRAM量が増えるようなグラフになった

        ![スクリーンショット 2023-08-27 17.57.13.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/144584/8c5030f5-8897-1115-dbf9-ce4cbefbff97.png)

    + 出力トークン数を64に固定し、入力トークン数を50刻みで増やして計測した結果はこちら


        ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/144584/341ba39b-8d43-9cf2-5de3-ec7daa682834.png)
        線形に増加しているというよりは、段々に増えているように見える


+ 出力トークンと生成時間は線形に増加していそう
    + `y = 0.47x + 9.08`だったので、傾き・切片共に大きめだった。