Vicuna v1.5をCTranslate2で変換して動かす on Databricks

CTranslate2関連で書きたいことがいくつかあるので、集中して書いていってみようと思います。
あと、今回まったくといっていいほどDatabricks関係ない。動作検証しただけで、タグ詐欺かもしれない。。。

CTranslate2とは？

CTranslate2は、Transformerモデルによる効率的な推論のためのC++とPythonのライブラリです。

このプロジェクトは、CPUとGPU上でTransformerモデルの高速化とメモリ使用量の削減を行うために、重みの量子化、レイヤーの融合、バッチの並べ替えなど、多くのパフォーマンス最適化技術を適用するカスタムランタイムを実装しています。

LLMを量子化するライブラリとしてはbitsandbytesやLlama.cpp、AutoGPTQなどが有名ですし、高速推論についてはHuggingface TGIやvLLMなどがありますが、CTranslate2も高速推論の機能を提供します。

以前、Rinna社の方が生成速度比較としていくつかのライブラリを比較していましたが、CTranslate2もそれに含まれており、悪くないパフォーマンスを示していました。

Databricks上で利用可能な高速推論ライブラリとしては、クセもなく個人的に使い勝手の良い優れた選択肢だと感じています。
今回は、このCTranslate2を使ってlmsys/vicuna-13b-v1.5を動かします。

Vicuna v1.5とは？

LMSysが開発したLLMで、Llama2ベースで構築されています。ライセンス体系もLlama2のライセンスに則っています。
transformersでの動作などは御大の記事を参照ください。

Llama2は日本語で回答を得るのにかなり苦労したのですが、Vicunaは日本語でも素直に回答を得やすいモデルだと思いますし、かなり優秀です。
今回は使いませんが、16K長のコンテキストサイズのモデルも公開されています。

というわけでやっていきます。

準備

Llama2をCTranslate2で使う方法はCTranslate2のgithub上にサンプルが公開されています。
ここを見れば、あとは読む必要がありません。

。。。というのは寂しいので、ちょっとだけ変更したやり方で進めます。

まずはhuggingface hubからlmsys/vicuna-13b-v1.5のモデルをダウンロードします。
モデル名が違うだけで、やり方は↓の記事参考。

また、pipで必要なモジュールをインストールしておきます。

%pip install -U -qq ctranslate2 sentencepiece transformers accelerate

dbutils.library.restartPython()

CTranslate2でモデル変換

上の記事とだいたい同じですが、ct2-transformers-converterでモデルを変換します。
githubのexampleだと--copy_filesで指定するのはtokenizer.modelだけですが、transformersのtokenizerを使えるようにするために、pecial_tokens_map.jsonとtokenizer_config.jsonも指定しておきます。
src_pathは変換元のモデルが格納されたパス、output_dirは変換済モデルの出力先パスです。

--quantizationにint8_float16を指定することで8bitに量子化します。

!ct2-transformers-converter --model {src_path} --copy_files tokenizer.model special_tokens_map.json tokenizer_config.json --output_dir {output_dir} --quantization int8_float16 --low_cpu_mem_usage

16Kコンテキストモデルは2023/8現在、まだ対応していないようです。
モデルの変換自体はエラーなくできるのですが、推論結果がおかしくなります。

以下のIssueに既に登録されていました。
https://github.com/OpenNMT/CTranslate2/issues/1403

推論

githubのexampleは少し複雑なやり方で実装されていますが、バッチ推論するだけならもう少し簡単に書けます。
また、exampleはTokenizerにSentencepieceをそのまま使っていましたが、汎用性を考えて今回はtransformersのAutoTokenizerでTokenizerを作成します。
そのために、変換時に必要なファイルをコピーしていました。

import torch
import transformers
import ctranslate2

model_path = "変換済みモデルのパス"

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"device: {device}")

generator = ctranslate2.Generator(model_path, device=device)
tokenizer = transformers.AutoTokenizer.from_pretrained(model_path)

# 関数でラップ
def generate_batch(instruction:str) -> str:

    prompt = f"""
A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions.

USER: {instruction}
ASSISTANT: """

    # 推論の実行
    tokens = tokenizer.convert_ids_to_tokens(
        tokenizer.encode(prompt, add_special_tokens=False)
    )
    results = generator.generate_batch(
        [tokens],
        max_length=256,
        sampling_topk=10,
        sampling_temperature=0.1,
        repetition_penalty=1.1,
        include_prompt_in_result=False,
        cache_static_prompt=True,
        return_scores=True,
    )

    text = tokenizer.decode(results[0].sequences_ids[0])
    return text

プロンプトはシステムプロンプト + USER: ASSISTANT: の指定で行います。

いくつか試してみます。

print(generate_batch("Hello!"))

結果

你好！有什么我可以帮助你的吗？

なぜか中国語で返してくる。

print(generate_batch("まどか☆マギカでは誰が一番かわいい？その理由も説明して。"))

結果

「魔法少女まどか☆マギカ」には、多くのかわいいキャラクターが登場しますが、一番かわいいと思うのは皆さんによって異なるでしょう。ただ、ここでは主人公の美樹さやかを例にして、そのかわいさについて説明してみます。

美樹さやかは、物語の開始時点では普通の高校生だったが、魔法戦争に巻き込まれ、魔法少女として覚醒します。彼女は、明るく前向きな性格で、周りの人々を和ませる存在です。また、魔法戦争においても、自分の命をかけて他の魔法少女を守ろうとする勇気と優しさがありま

御大のマネ。意外とまともな回答ですが、内容は間違ってますね。

タスクを実行させる

日本語知識はそこまで大したことがないので、単純な問いかけではVicunaの力を測ることはできません。

というわけで、In-context learningで要約・説明させてみます。

In-contextしない場合：

print(generate_batch("""織田信長がどういった人物だったかを40文字程度で説明しなさい。"""))

結果

織田信長は、戦国時代の日本で活躍した武将・大名であり、天下統一を目指す野心的な人物だった。

In-contextする場合：

# 織田信長のWikipediaの一部を使用
# https://ja.wikipedia.org/wiki/%E7%B9%94%E7%94%B0%E4%BF%A1%E9%95%B7

print(generate_batch("""
以下の文章から、織田信長がどういった人物だったかを40文字程度で説明しなさい。

織田信長は、織田弾正忠家の当主・織田信秀の子に生まれ、(中略) 近年の歴史学界では信長の評価の見直しが進んでいる[注釈 11]。
"""))

結果

織田信長は、戦国時代の日本を代表する武将・大名であり、尾張一国の支配者として勢力を拡大した。室町幕府再興を目指し、天下布武の印を使用したことで知られる。しかし、敵対勢力に追い詰められ、最後は家臣の明智光秀によって本能寺で自害に追い込まれた。信長は、極めて残虐な性格とされるが、政策面では協調的な姿勢を取り、勤王家として称賛されるようになった。

既に織田信長を知っているようなのであまりいい例ではないのですが、割と適切に回答してくれます。

というわけで、以下でされているようにQA等の利用で性能を発揮しそうです。

実行環境

ノードタイプ：g5.2xlargeで動作を確認しました。
上の推論実行後のVRAM使用量は13.67GB。
CTranslate2を使うことで、A10でもこの規模のモデルで高速に推論できるようになります。

まとめ

途中からVicunaの説明が中心でしたが、CTranslate2を使うと推論速度・メモリ使用量をバランスよく効率化できます。体感での性能劣化はそこまででもないと思います。

次回はCTranslate2で変換したモデルをlangchainで使う内容を投稿したいと思っています。