HuggingFaceのLLMのメモリ使用率を調べる

とあるプロジェクトでLLMを触っているが、LLMをはじめ言語のNNやHuggingFaceを触るのは初めてなので、そもそものリソース使用率を調べる

対象は日本語が使える rinna と open-calm

## 設定

環境はこんなところ

- GCE
- イメージ
    - `c0-deeplearning-common-cu113-v20230412-debian-10`
- n1-highmem-8
    - 8コア
    - 52GBメモリ
    - 200GBストレージ
- NVIDIA T4
    - 16GBメモリ

```Dockerfile
FROM python:3.10-slim

WORKDIR /app

COPY . .

RUN pip install \
    torch \
    transformers \
    sentencepiece \
    memory-profiler \
    accelerate \
    ;
```

```docker-compose.yaml
services:
  app:
    image: foo
    build:
      context: ./app
    deploy:
      resources:
        limits:
          cpus: '4.0'
          memory: 48G
        reservations:
          cpus: '2.0'
          memory: 20G
    volumes:
      - ./app:/app
    command: [ "python", "main.py", "***" ]  # 使うモデルを入れる
```

実行時引数でモデルを与えてメモリ使用率を見る

ややこしい文章課題は行わない
(長文の要約など)

## rinna

```main.py
import sys
import time
from logging import getLogger

import torch
from memory_profiler import profile
from transformers import AutoTokenizer, AutoModelForCausalLM

logger = getLogger(__name__)

dict_model = {
    "rinna": "rinna/japanese-gpt-neox-3.6b-instruction-sft-v2",
    "calm": "cyberagent/open-calm-1b",
}
prompter = {
    "rinna": lambda x: f"ユーザー: {x}<NL>システム: ",
    "calm": lambda x: x + ""
}

fout = open("musage.txt", "w")

@profile(stream=fout)
def main():
    start = time.time()
    logger.warning(f"start: {start - start}")

    try:
        framework = dict_model[sys.argv[1]]
    except Exception:
        logger.error("give an arg of rinna or calm")
        raise

    tokenizer = AutoTokenizer.from_pretrained(framework, use_fast=False)
    now = time.time()
    logger.warning(f"tokenizer loaded: {now - start}")

    model = AutoModelForCausalLM.from_pretrained(framework)
    now = time.time()
    logger.warning(f"model loaded: {now - start}")

    if torch.cuda.is_available():
        model = model.to("cuda")

    prompt = "フリーキャッシュフローとは何ですか"
    prompt = prompter[sys.argv[1]](prompt)
    logger.warning(f"prompt: {prompt}")

    token_ids = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")
    now = time.time()
    logger.warning(f"prompt encoded: {now - start}")

    with torch.no_grad():
        output_ids = model.generate(
            token_ids.to(model.device),
            do_sample=True,
            max_new_tokens=128,
            temperature=0.7,
            repetition_penalty=1.1,
            pad_token_id=tokenizer.pad_token_id,
            bos_token_id=tokenizer.bos_token_id,
            eos_token_id=tokenizer.eos_token_id
        )

    now = time.time()
    logger.warning(f"response generated: {now - start}")

    output = tokenizer.decode(output_ids.tolist()[0][token_ids.size(1):])
    now = time.time()
    logger.warning(f"token decoded: {now - start}")

    output = output.replace("<NL>", "\n")
    print(output)


if __name__ == "__main__":
    main()
```

モデルの取得に一番時間がかかる
回答の生成にも時間がかかる
cpu だけでも動く

```bash
temp-app-1  | start: 0.0
temp-app-1  | tokenizer loaded: 0.5645828247070312
temp-app-1  | model loaded: 35.02510094642639
temp-app-1  | prompt: ユーザー: フリーキャッシュフローとは何ですか<NL>システム:
temp-app-1  | prompt encoded: 35.029651165008545
temp-app-1  | response generated: 55.829381704330444
temp-app-1  | token decoded: 55.838337659835815
temp-app-1  | フリーキャッシュフローとは、現金の純増分のことです。企業が製品やサービスを販売し、顧客との取引から得た支払いを使っていくことで、お金が増えます。財務的な観点からは、こ
れは「良い」と言えます。

```

モデルのロードで14GB使っているが、実際には20GBは必要
ちなみに sft-v2 は事前学習済みのモデルを対話用にファインチューニングしたもの (なはず) で、
rinna の中では重い方

```python
Filename: /app/main.py

Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
    22    318.5 MiB    318.5 MiB           1   @profile(stream=fout)
    23                                         def main():
    24    318.5 MiB      0.0 MiB           1       start = time.time()
    25    318.5 MiB      0.0 MiB           1       logger.warning(f"start: {start - start}")
    26                                         
    27    318.5 MiB      0.0 MiB           1       try:
    28    318.5 MiB      0.0 MiB           1           framework = dict_model[sys.argv[1]]
    29                                             except Exception:
    30                                                 logger.error("give an arg of rinna or calm")
    31                                                 raise
    32                                         
    33    332.4 MiB     13.8 MiB           1       tokenizer = AutoTokenizer.from_pretrained(framework, use_fast=False)
    34    332.4 MiB      0.0 MiB           1       now = time.time()
    35    332.4 MiB      0.0 MiB           1       logger.warning(f"tokenizer loaded: {now - start}")
    36                                         
    37  14387.8 MiB  14055.5 MiB           1       model = AutoModelForCausalLM.from_pretrained(framework)
    38  14387.8 MiB      0.0 MiB           1       now = time.time()
    39  14387.8 MiB      0.0 MiB           1       logger.warning(f"model loaded: {now - start}")
    40                                         
    41  14388.0 MiB      0.1 MiB           1       if torch.cuda.is_available():
    42                                                 model = model.to("cuda")
    43                                         
    44  14388.0 MiB      0.0 MiB           1       prompt = "フリーキャッシュフローとは何ですか"
    45  14388.0 MiB      0.0 MiB           1       prompt = prompter[sys.argv[1]](prompt)
    46  14388.0 MiB      0.0 MiB           1       logger.warning(f"prompt: {prompt}")
    47                                         
    48  14388.0 MiB      0.0 MiB           1       token_ids = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")
    49  14388.0 MiB      0.0 MiB           1       now = time.time()
    50  14388.0 MiB      0.0 MiB           1       logger.warning(f"prompt encoded: {now - start}")
    51                                         
    52  14498.7 MiB      0.0 MiB           2       with torch.no_grad():
    53  14498.7 MiB    110.8 MiB           2           output_ids = model.generate(
    54  14388.0 MiB      0.0 MiB           1               token_ids.to(model.device),
    55  14388.0 MiB      0.0 MiB           1               do_sample=True,
    56  14388.0 MiB      0.0 MiB           1               max_new_tokens=128,
    57  14388.0 MiB      0.0 MiB           1               temperature=0.7,
    58  14388.0 MiB      0.0 MiB           1               repetition_penalty=1.1,
    59  14388.0 MiB      0.0 MiB           1               pad_token_id=tokenizer.pad_token_id,
    60  14388.0 MiB      0.0 MiB           1               bos_token_id=tokenizer.bos_token_id,
    61  14388.0 MiB      0.0 MiB           1               eos_token_id=tokenizer.eos_token_id
    62                                                 )
    63                                         
    64  14498.7 MiB      0.0 MiB           1       now = time.time()
    65  14498.7 MiB      0.0 MiB           1       logger.warning(f"response generated: {now - start}")
    66                                         
    67  14498.7 MiB      0.0 MiB           1       output = tokenizer.decode(output_ids.tolist()[0][token_ids.size(1):])
    68  14498.7 MiB      0.0 MiB           1       now = time.time()
    69  14498.7 MiB      0.0 MiB           1       logger.warning(f"token decoded: {now - start}")
    70                                         
    71  14498.7 MiB      0.0 MiB           1       output = output.replace("<NL>", "\n")
    72  14498.7 MiB      0.0 MiB           1       print(output)
```

## Open-CALM

1B で CALM の中では中くらいの重さ

rinna から細かいパラメータを変えた

```main.py
import sys
import time
from logging import getLogger

import torch
from memory_profiler import profile
from transformers import AutoTokenizer, AutoModelForCausalLM

logger = getLogger(__name__)

dict_model = {
    "rinna": "rinna/japanese-gpt-neox-3.6b-instruction-sft-v2",
    "calm": "cyberagent/open-calm-1b",
}
prompter = {
    "rinna": lambda x: f"ユーザー: {x}<NL>システム: ",
    "calm": lambda x: x,
}

fout = open("musage.txt", "w")

@profile(stream=fout)
def main():
    start = time.time()
    logger.warning(f"start: {start - start}")

    try:
        framework = dict_model[sys.argv[1]]
    except Exception:
        logger.error("give an arg of rinna or calm")
        raise

    tokenizer = AutoTokenizer.from_pretrained(framework)
    now = time.time()
    logger.warning(f"tokenizer loaded: {now - start}")

    model = AutoModelForCausalLM.from_pretrained(
        framework, device_map="auto", torch_dtype=torch.float32,
    )
    now = time.time()
    logger.warning(f"model loaded: {now - start}")

    if torch.cuda.is_available():
        model = model.to("cuda")

    prompt = "フリーキャッシュフローとは何ですか"
    prompt = prompter[sys.argv[1]](prompt)
    logger.warning(f"prompt: {prompt}")

    token_ids = tokenizer.encode(prompt, return_tensors="pt")
    now = time.time()
    logger.warning(f"prompt encoded: {now - start}")

    with torch.no_grad():
        output_ids = model.generate(
            token_ids.to(model.device),
            do_sample=True,
            max_new_tokens=128,
            temperature=0.7,
            top_p=0.9,
            repetition_penalty=1.1,
            pad_token_id=tokenizer.pad_token_id,
            #bos_token_id=tokenizer.bos_token_id,
            #eos_token_id=tokenizer.eos_token_id
        )

    now = time.time()
    logger.warning(f"response generated: {now - start}")

    output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    now = time.time()
    logger.warning(f"token decoded: {now - start}")

    output = output.replace("<NL>", "\n")
    print(output)


if __name__ == "__main__":
    main()
```

huggingfaceのドキュメント通りに利用しようとするとエラーが出る
`RuntimeError: "LayerNormKernelImpl" not implemented for 'Half'`
https://github.com/huggingface/transformers/issues/21989#issuecomment-1461999591
float16を使うにはGPUが必要らしい
なので `torch_ftype=torch.float32` を使った

とりあえず CPU だけでも動く

モデルのロードは比較的早い
それと比較すると回答の生成に時間がかかる上に、回答が意図したものではない

```
temp-app-1  | start: 0.0
temp-app-1  | tokenizer loaded: 0.3371620178222656
temp-app-1  | The model weights are not tied. Please use the `tie_weights` method before using the `infer_auto_device` function.
temp-app-1  | model loaded: 8.491863489151001
temp-app-1  | prompt: フリーキャッシュフローとは何ですか
temp-app-1  | prompt encoded: 8.495642185211182
temp-app-1  | response generated: 36.54288625717163
temp-app-1  | token decoded: 36.544095516204834
temp-app-1  | フリーキャッシュフローとは何ですか。
temp-app-1  | 【財務】負債比率とはどういう意味でしょうか?また、その計算方法を教えてください。【ESG/CSR】【労務管理・衛生要因分析(HACCP)】「経営戦略」と「組織運営方針」「人事制度等 の基本的仕組み」、これら3つについて教えてください。(製造業)【経理業務効率化支援ツール】(経費精算システム)「勤怠管理システム」、「会計帳簿作成ソフト」)のそれぞれの機能概要はどのようなものなのでしょうか?【総務系サービス開発プロジェクト事例紹介】・社内イントラネットに社員ブログを設置したい。・社内で使用するPCが老朽化したため新しく買い替えたい
```

モデルサイズも rinna よりは小さいのか、メモリを6GBほどしか使わない

```python
Filename: /app/main.py

Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
    22    314.6 MiB    314.6 MiB           1   @profile(stream=fout)
    23                                         def main():
    24    314.6 MiB      0.0 MiB           1       start = time.time()
    25    314.6 MiB      0.0 MiB           1       logger.warning(f"start: {start - start}")
    26                                         
    27    314.6 MiB      0.0 MiB           1       try:
    28    314.6 MiB      0.0 MiB           1           framework = dict_model[sys.argv[1]]
    29                                             except Exception:
    30                                                 logger.error("give an arg of rinna or calm")
    31                                                 raise
    32                                         
    33    332.8 MiB     18.2 MiB           1       tokenizer = AutoTokenizer.from_pretrained(framework)
    34    332.8 MiB      0.0 MiB           1       now = time.time()
    35    332.8 MiB      0.0 MiB           1       logger.warning(f"tokenizer loaded: {now - start}")
    36                                         
    37   6930.0 MiB   6597.2 MiB           2       model = AutoModelForCausalLM.from_pretrained(
    38    332.8 MiB      0.0 MiB           1           framework, device_map="auto", torch_dtype=torch.float32,
    39                                             )
    40   6930.0 MiB      0.0 MiB           1       now = time.time()
    41   6930.0 MiB      0.0 MiB           1       logger.warning(f"model loaded: {now - start}")
    42                                         
    43   6930.0 MiB      0.0 MiB           1       if torch.cuda.is_available():
    44                                                 model = model.to("cuda")
    45                                         
    46   6930.0 MiB      0.0 MiB           1       prompt = "フリーキャッシュフローとは何ですか"
    47   6930.0 MiB      0.0 MiB           1       prompt = prompter[sys.argv[1]](prompt)
    48   6930.0 MiB      0.0 MiB           1       logger.warning(f"prompt: {prompt}")
    49                                         
    50   6930.0 MiB      0.0 MiB           1       token_ids = tokenizer.encode(prompt, return_tensors="pt")
    51   6930.0 MiB      0.0 MiB           1       now = time.time()
    52   6930.0 MiB      0.0 MiB           1       logger.warning(f"prompt encoded: {now - start}")
    53                                         
    54   6957.4 MiB      0.0 MiB           2       with torch.no_grad():
    55   6957.4 MiB     27.4 MiB           2           output_ids = model.generate(
    56   6930.0 MiB      0.0 MiB           1               token_ids.to(model.device),
    57   6930.0 MiB      0.0 MiB           1               do_sample=True,
    58   6930.0 MiB      0.0 MiB           1               max_new_tokens=128,
    59   6930.0 MiB      0.0 MiB           1               temperature=0.7,
    60   6930.0 MiB      0.0 MiB           1               top_p=0.9,
    61   6930.0 MiB      0.0 MiB           1               repetition_penalty=1.1,
    62   6930.0 MiB      0.0 MiB           1               pad_token_id=tokenizer.pad_token_id,
    63                                                     #bos_token_id=tokenizer.bos_token_id,
    64                                                     #eos_token_id=tokenizer.eos_token_id
    65                                                 )
    66                                         
    67   6957.4 MiB      0.0 MiB           1       now = time.time()
    68   6957.4 MiB      0.0 MiB           1       logger.warning(f"response generated: {now - start}")
    69                                         
    70   6957.4 MiB      0.0 MiB           1       output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    71   6957.4 MiB      0.0 MiB           1       now = time.time()
    72   6957.4 MiB      0.0 MiB           1       logger.warning(f"token decoded: {now - start}")
    73                                         
    74   6957.4 MiB      0.0 MiB           1       output = output.replace("<NL>", "\n")
    75   6957.4 MiB      0.0 MiB           1       print(output)
```

# 終わりに

正直、CPUだけでも動くか、メモリ使用率の目安を知りたかった
推論のためだけにGPUを使うのはコストなので、CPUだけでも動くようでよかった
まぁ、大きなモデルを高速に動かそうと思ったらGPUいるかもだが