tinyrwkv

現在、 rewriteブランチで書き直されています。

tinyrwkv: RWKV-LM のより小さな移植版

大規模言語モデルの RWKV-LM ファミリを tinygrad フレームワークに移植したもの。

ロードマップ

• WKV カーネルをカスタム関数として実装する
• WKV カーネルの後方をカスタム関数として実装する
• ワールドモデルとトークナイザーのサポートを追加
• MIDIモデルのサポートを追加
• RWKV-5 モデルの初期サポートを追加

依存関係

現在、git から tinygrad を使用するか、nix flake を使用する必要があります。

パイソン

 numpy
pydot (only for GRAPH=1)
tinygrad
tokenizers
torch (only for loading pytorch weights)
tqdm
wandb (optional during training)

システム

 rust (only for compiling)
clang (only for compiling)
graphviz (only for GRAPH=1)

使用法

python -m cliを使用して CLI を実行します。

また、他のプロジェクトに埋め込むための Python パッケージとしても使用できます。モデルを移植可能な C コードにコンパイルし、そのように埋め込むことも可能です。

 usage: tinyrwkv-cli [-h] [--seed SEED] {pre,gen,cht,cmp,bch,ptr,gpt,tra,bpt,wkv,mus} ...

CLI for tinyrwkv

positional arguments:
  {pre,gen,cht,cmp,bch,ptr,gpt,tra,bpt,wkv,mus}
    pre                 preprocess either tinyrwkv trained weights or pytorch trained weights into RNN form
    gen                 freeform generation using the RNN mode (requires a preprocessed model using `pre`)
    cht                 chat with a model in RNN mode (requires a preprocessed model using `pre`)
    cmp                 compile a RNN model into c source code and a compiled executable (need to run with CLANG=1)
    bch                 benchmark the rnn mode
    ptr                 preprocess pytorch weights weights into GPT form for training or inference
    gpt                 freeform generation using the GPT mode (requires a preprocessed model using `ptr`)
    tra                 pretrain or finetune a model (if finetuning the model needs to be preprocessed with `ptr`)
    bpt                 benchmark the gpt mode
    wkv                 benchmark/test each wkv module
    mus                 music generation using the RNN mode (requires a preprocessed model using `pre`)

options:
  -h, --help            show this help message and exit
  --seed SEED           seed for random

ライセンス

LICENSE ファイルと NOTICE ファイルを参照してください。