自宅マシンでFugaku-LLM-13B を動かすために(複数GPU)

Object

• 自宅マシンでFugaku-LLMをGPUで動かす環境を作る
• 13BのLLMモデルを自宅用GPUで乗せた時のメモリ使用量実績の記事があまりなかったので、メモ
• Fugaku-LLMならではの話は触れていないので、別記事で中身の検証色々やっていく

Summary

• Fugaku-LLMは13Bのパラメータを持つ。そのため、家庭用GPU1枚の上限24GBメモリに乗りきらない。(約26~28GB)
• メモリ利用量:そのため、win10上のRTX3090(24GB),RTX3070(8GB)の2台でロード、計約28.6GB。GPUそれぞれの利用量は以下参照
  • RTX3090:21.6GB/24GB (MSI Ventus 24G)
    • メインGPUなのでwinで約600MB利用していたのを含む
    • ライザーケーブル利用
  • RTX3070:7GB/8GB (Palit Gaming Pro 8G)
• 　速度:LLMChain関数を用いて、llm_chain.runを実行したところ、10s以内には大体応答帰ってくる。十分利用できる
  • 　実行プロンプトは、公式READMEを参考に以下のような形

以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。

### 指示:
おいしいカレーを作る方法を教えてください。

### 応答:

• device_map=autoだと、cpuへのオフロードが行われたため、device_mapで詳細にデバイス指定を行った

メモ

• device_map=autoの時に以下Warningがモデルロード時に出る。DRAMを使っているらしく、極端に推論速度が落ちたので対策必要だった。

  WARNING:root:Some parameters are on the meta device device because they were offloaded to the cpu.

• そのため、以下のようなコードで一旦Warning回避

for i in range(30):
    device_map["transformer.h." + str(i)] = 0
for i in range(30, 100):
    device_map["transformer.h." + str(i)] = 1
    # 実際にはtransformer.h.35~40くらいまでしかないはず
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    low_cpu_mem_usage=True,
    device_map=device_map,
    torch_dtype=torch.bfloat16,
) 

• 全体のコードは以下

import os
from langchain import HuggingFacePipeline, PromptTemplate, LLMChain
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

model_name = "Fugaku-LLM/Fugaku-LLM-13B-instruct"
task = "text-generation"
device_map = {
    "transformer.wte": 0,
    "transformer.wpe": 0,
    "transformer.ln_f": 0,
    "lm_head": 0,
}

for i in range(30):
    device_map["transformer.h." + str(i)] = 0
for i in range(30, 100):
    device_map["transformer.h." + str(i)] = 1
    # 実際にはtransformer.h.35~40くらいまでしかないはず
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    low_cpu_mem_usage=True,
    device_map=device_map,
    torch_dtype=torch.bfloat16,
)
pipe = pipeline(
    task,
    model=model,
    tokenizer=tokenizer,
    framework="pt",
    max_new_tokens=1024,
)
llm = HuggingFacePipeline(pipeline=pipe)
template = "{system_example}\n\n### 指示:\n{instruction_example}\n\n### 応答:\n"
prompt = PromptTemplate.from_template(template)
system_example = (
    "以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。"
)
instruction_example = "おいしいカレーを作る方法を教えてください。"
llm_chain = LLMChain(prompt=prompt, llm=llm, verbose=True)
response = llm_chain.run(
    {"system_example": system_example, "instruction_example": instruction_example}
)
print(response)

• なお、回答は以下の通り。さて・・・おいしいカレーはできるのか。やってみる。

以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。

### 指示:
おいしいカレーを作る方法を教えてください。

### 応答:
おいしいカレーを作るには、いくつかの重要なステップがある。まず、肉を焼き、次に玉ねぎを炒め、次にカレーパウダーを加え、次に他のスパイスを加え、最後に水を加え、すべてのスパイスをよく混ぜ合わせ、次にトマトペーストを加え、すべてのスパイスをよく混ぜ合わせ、最後に他のスパイスを加え、すべてのスパイスをよく混ぜ合わせ、最後に1時間ほど煮込む。