Stability AIの日本語モデルをRTX3060(12GB)シングル環境で使う方法

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

過去のエントリ「LLaMA 2モデルをRTX3060(12GB)シングル環境で使う方法」に続き、今回もオンプレで商用でも使えるというStability AIの日本語モデルJapanese StableLM AlphaをRTX3060(12GB)環境で動かしてみたので、その備忘録。

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

対応

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

※私の検証環境ではそのまま動かすとVRAM不足によりエラーとなってしまうため、8ビット量子化して動かしています(8ビット量子化でのVRAM使用量は約9GB)。

サンプルソースコード

pip:

accelerate
bitsandbytes
einops
scipy
sentencepiece
torch
transformers

main.py

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

MODEL_ID = 'stabilityai/japanese-stablelm-base-alpha-7b'

DEVICE = 'cuda'

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 = LlamaTokenizer.from_pretrained(
    'novelai/nerdstash-tokenizer-v1',
    additional_special_tokens=['▁▁'],
    legacy=False
)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
    quantization_config=config,
    pad_token_id=3
)

model.eval()

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

    outputs = model.generate(
        inputs,
        max_new_tokens=128,
        temperature=0.5,
        top_p=0.95,
        do_sample=True,
    )

    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.')

Hugging Faceに公開されているサンプルソースをベースに若干手を入れています。

実行結果

プロンプト(1):

私はアリス。

生成テキスト(1):

アリス・フォン・アルノルト。
アルノルト侯爵家の次女で、今年16歳になる。
私の婚約者は、この国の王太子、レオナルド・フォン・アルノルト。
この国の第一王子で、私より4つ歳上の19歳。
そして、私の元婚約者は、この国の第一王子、レオナルド・フォン・アルノルト。
この国の第一王子で、私より4つ歳上の19歳。
そして、私の元婚約者は、この国の第一王子、レオナルド・フォン・アルノルト。
この国の第一王子で、私より4つ歳上の19

プロンプト(2):

「Stability AIの日本語モデルをRTX3060(12GB)シングル環境で使う方法」を英訳すると

生成テキスト(2):

「Stability AIの日本語モデルをRTX3060(12GB)シングル環境で使う方法」になるように、英訳を工夫してみてください。
英訳のポイント
英訳のポイントは、日本語の文章をそのまま英訳しないことです。日本語の文章をそのまま英訳してしまうと、文法や単語の使い方がおかしくなってしまいます。
日本語の文章をそのまま英訳するのではなく、英訳しやすいように日本語の文章を工夫しましょう。
日本語の文章を工夫する方法
日本語の文章を工夫する方法は、英語に翻訳しやすいように、文章の順番を入れ

8ビット量子化した状態で使ってみた感想としては、意味不明な回答であったりオウム返しになったり、過学習っぽい応答があるなど、何となくファインチューニングしないと素の状態では使えそうにないという印象です。パラメータ調整で多少は改善できるのかもしれませんが…。

ちなみにChatGPTのように質問に対する回答が生成できるJapanese-StableLM-Instruct-Alpha-7Bについては、研究目的での利用限定で商用利用はNGということに注意が必要です。
(利用には事前の承認申請とHugging Face Hubへのログインが必要。)

参考ウェブサイトなど

• Hugging Face
  Stability AI

以上です。