OllamaでLlama3.1を使ったときは速いくせに、huggingfaceでLlama3.1を使ったときは遅いときの解決をしようとしたができなかった話

## はじめに

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

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

```python
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 のほうが若干速い、程度のはず、、

しかし、実行結果を見てみると、、

![推論時間](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3694964/362cb75c-c79c-ceaa-a3c4-1e53d3ef122c.png)

14秒ほどの差があります。

## 解決への試み

1. パディングの追加
   ```python
   model_input = tokenizer(prompt, padding=True, truncation=True, max_length=1024, return_tensors="pt").to(model.device)
   ```
   結果: 大きな改善は見られませんでした。

2. 二重量子化の無効化
   ```python
   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秒短縮されました。

   ![改善後の推論時間](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3694964/e1eddcaa-c783-1de8-b2a5-d8ce4ba149c2.png)

## メモリ使用量の比較

- 二重量子化有効時:
  ![True時のメモリ使用量](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3694964/b91a4a3a-c8d9-3c05-8c3b-c4ea2635f33a.png)

- 二重量子化無効時:
  ![False時のメモリ使用量](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3694964/abdc8520-e62a-6ef4-b5f3-a689d03cc14d.png)

メモリ使用量の差はわずかですが、推論時間に大きな影響がありました。

## FlashAttention-2の導入

さらなる最適化を目指し、FlashAttention-2を導入しました。これはよりメモリ効率の良いアテンションメカニズムです。

### 準備作業

1. CUDA Toolkitのインストール (WSL環境用)

CUDA Tool kitをインストールする必要があるのでインストールします。[こちら](https://developer.nvidia.com/cuda-downloads?target_os=Linux&target_arch=x86_64&Distribution=Ubuntu&target_version=24.04&target_type=deb_network)

これらのコードで自分の環境に合うCUDAをインストールしてください。

```bash
cat /etc/os-release
uname -m
```

実行した結果以下のような選択をしました。wslを使っている環境の人はこれでいいと思います。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3694964/92869f97-7b26-b5aa-1368-ceb8d7d85b5f.png)

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

2. 環境変数の設定
   ```bash
   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}}
   ```

3. PyTorchの再インストール (CUDAにあったバージョンで)
   ```bash
   pip uninstall torch torchvision
   pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121
   ```

4. 必要なパッケージのインストール
   ```bash
   pip install packaging ninja
   ```

### FlashAttention-2の実装

モデルの初期化時に以下のパラメータを追加:

```python
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    quantization_config=bnb_config,
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2"
)
```

### 結果

1. パディングありの場合:
   ![FlashAttention-2結果(パディングあり)](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3694964/49554b12-d8cf-4eb3-87a4-4a4d73d88639.png)
   
   予想に反して、推論時間が0.5秒ほど増加しました。

2. パディングなしの場合:
   ```python
   model_input = tokenizer(prompt, return_tensors="pt").to(model.device)
   ```
   ![FlashAttention-2結果(パディングなし)](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3694964/ab8102c4-fec1-f125-8527-1aa5df85f7fb.png)
   
   さらに推論時間が2秒ほど増加しました。これは予想外の結果でした。

## 結論

1. 二重量子化の無効化は推論時間の短縮に効果がありましたが、Ollamaモデルの性能には及びませんでした。

2. FlashAttention-2の導入は、予想に反して性能を低下させました。

いったんはOllamaを使おうと思います