Google Colab で RWKV を試す - RWKVパッケージ版
「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) に配置。
(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)
