LLaMA 2モデルをRTX3060(12GB)シングル環境で使う方法

この記事は公開から1年以上経過しています。

オンプレで商用でも使えてChatGPT(GPT3.5)にも匹敵すると言われているMetaの大規模言語モデルLLaMA 2が、どの程度なのか気になり早速ローカルのNVIDIA RTX3060(12GB)シングル環境で試してみたところ、過去のエントリ「Stable Diffusion(PyTorch)でOutOfMemoryが出るときの解決策」のVRAM断片化対策を行ってもVRAM不足で実行できなかったため、そのときの対応についての備忘録。

検証環境はLinux(Ubuntu 22)＋GPU(RTX3060 12GB)＋Python(3.10)です。

2023.8.15追記:
8ビット量子化でもVRAM使用量8GB程度で実行可能でしたので8ビット量子化についての説明を追記し、ソースも若干修正しました。

対応

VRAMサイズを抑えるために4ビット/8ビット量子化を行う。

※私の検証環境ではMeta公式の一番小さな7Bモデルでも、そのまま動かすとVRAM不足によりエラーとなってしまうため、ひとまず7Bを4ビット(NF4)/8ビット量子化して動かしています。

Hugging FaceのLLaMA 2モデルを利用するには事前にMeta社の申請フォームから申請を行い承認されている必要があります。

Hugging Face Hubのキャッシュディレクトリを変更する場合は過去のエントリ「Hugging Face Hubのキャッシュディレクトリを変更する方法」をご参照ください。

モデルのダウンロードサイズが7Bの場合でも14GB程度あります。通信環境にご注意ください。

サンプルソースコード

pip:

accelerate
bitsandbytes
scipy
torch
transformers
huggingface-hub

main.py:

import torch
from huggingface_hub import login
from transformers import (AutoModelForCausalLM, AutoTokenizer,
                          BitsAndBytesConfig)

DEVICE = 'cuda'

login(token='Hugging Face Auth Token Id')

# MODEL_ID = 'meta-llama/Llama-2-7b-hf'
MODEL_ID = 'meta-llama/Llama-2-7b-chat-hf'
# MODEL_ID = 'meta-llama/Llama-2-13b-hf'

config = None
if DEVICE == 'cuda':
    # 4ビット量子化の場合はこちら
    config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_use_double_quant=True,
        bnb_4bit_compute_dtype=torch.bfloat16
    )

    # # 8ビット量子化の場合はこちら
    # config = BitsAndBytesConfig(
    #    load_in_8bit=True
    # )

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, quantization_config=config
)

def get_model_response(prompt):
    inputs = tokenizer.encode(
        prompt,
        add_special_tokens=True,
        return_tensors='pt'
    ).to(DEVICE)

    outputs = model.generate(
        inputs,
        max_length=256,
        #  do_sample=True,
        top_k=50,
        top_p=0.95,
        temperature=0.5,
        num_return_sequences=1
    )

    response = tokenizer.decode(
        outputs[0][len(inputs[0]):], skip_special_tokens=True)
    return response

while True:
    print("Enter your text (or 'quit' to exit):")
    prompt = input().strip()
    if prompt.lower() == 'quit':
        break
    elif len(prompt) > 0:
        print(get_model_response(prompt))
    else:
        print('Nothing has been entered.')

本サンプルではBitsAndBytesConfigを直接使用していますが、NF4など細かい量子化設定を行わないのであればAutoModelForCausalLM.from_pretrained()のload_in_4bit=Trueやload_in_8bit=Trueで良さそうです。

ソース中のHugging Face Auth Token Idは、Hugging FaceのUser Profile→Settings→Access Tokensにあるご自身のgit cloneのトークンIDを使用します。Cliからログインする場合はlogin()部分は不要です。

実行結果

このモデルの実行結果(4bit量子化)サンプルは以下のとおりです。

プロンプト:

Q:Who are you?\nA:

回答:

I am LLaMA, an AI assistant developed by Meta AI that can understand and respond to human input in a conversational manner. nobody knows the answer to this question.

回答(和訳):

私はLLaMAです。Meta AIが開発したAIアシスタントで、人間の入力を理解し、会話形式で応答することができます。

補完タスクで一応それっぽい回答を返してきているようです。
現状はひとまず動くところまでしか試していませんが、もう少しパラメーターやプロンプトの調整が必要そうです。

ちなみに4bit/8bit量子化を行うとかなり精度がおちると言われていますので、モデルの精度を検証したいのであればGoogle Colaboratoryの有償版などを使ったほうが良さそうです。

参考ウェブサイトなど

Hugging Face
Meta Llama 2
Hugging Face
Making LLMs even more accessible with bitsandbytes, 4-bit quantization and QLoRA
Hugging Face
Quantize 🤗 Transformers models

以上です。