はじめに
Ollamaでllama3.1:8b-instruct-q3_K_Sを使用すると約10秒で推論できるのに対し、HuggingFaceからmeta-llama/Meta-Llama-3-8B-Instructを使用すると非常に遅くなる問題について検討します。
使用環境
- GPU: RTX 3060
- モデル: meta-llama/Meta-Llama-3.1-8B-Instruct
一番早かったコード
init.py
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
# グローバル変数を最初に宣言
model = None
tokenizer = None
model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct"
#tokenizerとmodelを作成、またはロード
def make_or_load_model_and_tokenizer():
global model, tokenizer
#作成する場合（一回目）
if model is None or tokenizer is None:
print("Loading model and tokenizer")
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.float16,
bnb_4bit_quant_type="nf4",
bnb_4bit_use_double_quant=False,
)
model = AutoModelForCausalLM.from_pretrained(
model_name,
device_map="auto",
quantization_config=bnb_config,
torch_dtype=torch.bfloat16
)
print("Model and tokenizer loaded")
else:
#作成していた場合（1<N回目）
print("Model and tokenizer already loaded")
return model, tokenizer
# モデルとトークナイザーを初期化
model, tokenizer = make_or_load_model_and_tokenizer()
default_use.py
import torch
import time
import os
import psutil
from init import model, tokenizer
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, pipeline
#text作成
def generate_text(prompt):
#truncationはmax_lengthに切ってくれる。
model_input = tokenizer(prompt,padding=True, truncation=True, max_length=4086, return_tensors="pt").to(model.device)
input_ids = model_input["input_ids"]
#勾配計算はする必要がないためno_gradを使用
# 開始時間を記録
start_time = time.time()
with torch.no_grad():
result = model.generate(
input_ids,
max_length=400,
do_sample=True,
temperature=0.6,
top_p=0.9
)
result = result[0][input_ids.shape[-1]:]
output = tokenizer.decode(result, skip_special_tokens=True)
print("\n生成結果\n",output)
# 終了時間を記録
end_time = time.time()
# 実行時間を計算
execution_time = end_time - start_time
print(f"Execution time: {execution_time:.2f} seconds")
del input_ids
del model_input
torch.cuda.empty_cache()
if __name__ == "__main__":
prompt = "Explain about LLM briefly."
generate_text(prompt)
問題点
HuggingFaceのモデルを使用すると、推論時間が非常に遅くなります。
ollamaのq3-K-Sとの比較:
推論速度はlama3.1:8b-instruct-q3_K_S のほうが若干速い、程度のはず、、
しかし、実行結果を見てみると、、
14秒ほどの差があります。
解決への試み
-
パディングの追加
model_input = tokenizer(prompt, padding=True, truncation=True, max_length=1024, return_tensors="pt").to(model.device)結果: 大きな改善は見られませんでした。
-
二重量子化の無効化
bnb_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=False )結果: 推論時間が約6秒短縮されました。
メモリ使用量の比較
-
二重量子化有効時:
-
二重量子化無効時:
メモリ使用量の差はわずかですが、推論時間に大きな影響がありました。
FlashAttention-2の導入
さらなる最適化を目指し、FlashAttention-2を導入しました。これはよりメモリ効率の良いアテンションメカニズムです。
準備作業
- CUDA Toolkitのインストール (WSL環境用)
CUDA Tool kitをインストールする必要があるのでインストールします。こちら
これらのコードで自分の環境に合うCUDAをインストールしてください。
cat /etc/os-release
uname -m
実行した結果以下のような選択をしました。wslを使っている環境の人はこれでいいと思います。
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb
sudo apt-get update
sudo apt-get -y install cuda-toolkit-12-6
-
環境変数の設定
export PATH=/usr/local/cuda-12.6/bin${PATH:+:${PATH}} export LD_LIBRARY_PATH=/usr/local/cuda-12.6/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}} -
PyTorchの再インストール (CUDAにあったバージョンで)
pip uninstall torch torchvision pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121 -
必要なパッケージのインストール
pip install packaging ninja
FlashAttention-2の実装
モデルの初期化時に以下のパラメータを追加:
model = AutoModelForCausalLM.from_pretrained(
model_name,
device_map="auto",
quantization_config=bnb_config,
torch_dtype=torch.bfloat16,
attn_implementation="flash_attention_2"
)
結果
-
パディングありの場合:
予想に反して、推論時間が0.5秒ほど増加しました。
-
パディングなしの場合:
model_input = tokenizer(prompt, return_tensors="pt").to(model.device)
さらに推論時間が2秒ほど増加しました。これは予想外の結果でした。
結論
-
二重量子化の無効化は推論時間の短縮に効果がありましたが、Ollamaモデルの性能には及びませんでした。
-
FlashAttention-2の導入は、予想に反して性能を低下させました。
いったんはOllamaを使おうと思います