Google Colab で RWKV を試す - RWKVパッケージ版

npaka

2023年4月9日 11:05

「Google Colab」で「RWKV」を試したので、まとめました。

・rwkv 0.7.3

前回

1. RWKV

「RWKV」は、TransformerレベルのLLM性能を備えたRNNです。高性能、高速推論、VRAMの節約、高速学習、長い文脈長、自由な埋め込みを実現しています。

2. Colabでの実行

Colabでの実行手順は、次のとおりです。

(1) メニュー「編集→ノートブックの設定」で、「ハードウェアアクセラレータ」で「GPU」を選択。

(2) Googleドライブのマウント。

# Googleドライブのマウント
from google.colab import drive
drive.mount('/content/drive')

(3) 作業フォルダへの移動。

# 作業フォルダへの移動
import os
os.makedirs("/content/drive/My Drive/work", exist_ok=True)
%cd '/content/drive/My Drive/work'

(4) モデルをダウンロードして、作業フォルダ (work) に配置。
RWKVのモデルとして、Pileのみ学習させたベースモデルとInstructionチューニングしたRavenがあります。Ravenはバージョン (v1〜v9…) と言語種別 (Eng、EngAndMore、EngChnJpn、ChnEngなど) もあります。

・Pile (ベースモデル)
　・rwkv-4-pile-14b
　・rwkv-4-pile-7b
　・rwkv-4-pile-3b
　・rwkv-4-pile-1b5
　・rwkv-4-pile-430m
　・rwkv-4-pile-169m

・Raven (Instructionチューニングモデル)
　・RWKV-4-Raven-14B-v<n>-XXX
　・RWKV-4-Raven-7B-v<n>-XXX
　・RWKV-4-Raven-3B-v<n>-XXX
　・RWKV-4-Raven-1B5-v<n>-XXX

今回は、「RWKV-4-Raven-7B-v7-EngAndMore-20230404-ctx4096.pth」を使いました。

(5) トークンファイルをダウンロードして、作業フォルダ (work) に配置。

・20B_tokenizer.json

(6) パッケージのインストール。

# パッケージのインストール
!pip install rwkv

(7) 環境変数の準備。

# 環境変数の準備
import os
os.environ['RWKV_JIT_ON'] = '1'
os.environ["RWKV_CUDA_ON"] = '0'

「RWKV_CUDA_ON」は、C++コンパイラとCUDAが必要ですが、CUDA カーネルをコンパイルして10倍高速にします。

(8) モデルとパイプラインの準備。

from rwkv.model import RWKV
from rwkv.utils import PIPELINE, PIPELINE_ARGS

# モデルとパイプラインの準備
model = RWKV(
    model="./RWKV-4-Raven-7B-v7-EngAndMore-20230404-ctx4096", 
    strategy="cuda fp16")
pipeline = PIPELINE(model, "./20B_tokenizer.json")

◎ model
モデルパスを指定します。セパレータには「\」ではなく「/」を使用します。長いctxが必要な場合は、「ctx4096」モデルを使用します。

◎ strategy
cpu or cuda と精度を指定します。

・fp16 = GPU に良い (CPU非サポート)
・fp32 = CPU に良い
・bf16 = 精度が悪い (CPUサポート)

設定例は、次のとおりです。

・'cpu fp32' = 全レイヤー cpu fp32
・'cuda fp16' = 全レイヤー cuda fp16
・'cuda fp16i8' = 全レイヤー int8 量子化によるcuda fp16
・'cuda fp16i8 *10 -> cpu fp32' = 最初の10レイヤー cuda fp16i8、次のレイヤー cpu fp32 (速度を上げるには10を増やす)
・'cuda:0 fp16 *10 -> cuda:1 fp16 *8 -> cpu fp32' = 最初の10レイヤー cuda:0 fp16、次の8レイヤー cuda:1 fp16、次のレイヤー cpu fp32

(9) パイプライン引数の準備。

# パイプライン引数の準備
args = PIPELINE_ARGS(
    temperature = 1.0,
    top_p = 0.7, 
    top_k = 100, 
    alpha_frequency = 0.25, 
    alpha_presence = 0.25, 
    token_ban = [],
    token_stop = [0],
    chunk_len = 256)

・temperature : 温度
・top_p : top-P
・top_k : top-K (0は無視)
・alpha_frequency : 同じテキスト出力をへらすため出力トークンに課すペナルティ (-2〜2)
・alpha_presence : 同じトークンの繰り返しを減らすため出力トークンに課すペナルティ (-2〜2)
・token_ban : 禁止トークン
・token_stop : 停止トークン
・chunk_len : チャンク長。入力を分割してVRAM節約で、短いほうが遅い。

・frequency_penalty : 出力トークンに課すペナルティ。同じテキスト出力の可能性を減らす（デフォルト:0, -2〜2）
・presence_penalty : 出力トークンに課すペナルティ。同じトークンの繰り返し可能性を減らす（デフォルト:0, -2〜2）

「alpha_frequency」「alpha_presence」について詳しくは、以下を参照。

(9) Instructプロンプトを生成する関数の準備。
「ChatRWKV」の「+i」と同じ処理になります。

# Instructプロンプトの生成
def generate_prompt(instruction, input=None):
    if input:
        return f"""Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.

# Instruction:
{instruction}

# Input:
{input}

# Response:
"""
    else:
        return f"""Below is an instruction that describes a task. Write a response that appropriately completes the request.

# Instruction:
{instruction}

# Response:
"""

(10) パイプラインの実行。

# Instructプロンプトの生成
prompt = generate_prompt("日本で一番人気のアニメは？")
print("--[prompt]--\n" + prompt + "----")

# パイプラインの実行
result = pipeline.generate(prompt, token_count=200, args=args)
print(result)

【おまけ】　ストリーミング出力

pipeline.generate()のcallbackでストリーミング出力も可能です。

# パイプラインの実行
def my_print(s):
    print(s, end='', flush=True)
pipeline.generate(prompt, token_count=200, args=args, callback=my_print)

次回

この記事が気に入ったらサポートをしてみませんか？

Google Colab で RWKV を試す - RWKVパッケージ版

前回

1. RWKV

2. Colabでの実行

【おまけ】 ストリーミング出力

関連

次回

【おまけ】　ストリーミング出力