はじめに

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を使おうと思います