Repolode
投稿する
BlinkDL

BlinkDL/RWKV-LM

GPTのTransformerのように並列学習できる、RNN型言語モデルアーキテクチャRWKVの公式学習・リファレンスコード。

GitHub で見る
スター
14,757
フォーク
1,025
未解決の Issue
159
最終コミット
23 日前

これは何？

RWKV-LMは、再帰型ニューラルネットワークでありながらTransformerのように並列学習できる言語モデルアーキテクチャRWKVの、メインとなる学習・リファレンスリポジトリです。現在のプロジェクトはRWKV-7「Goose」で、READMEでは、Attentionを使わず、時間計算量が線形で、メモリ使用量が一定、KV cacheも不要な設計として説明されています。推論時にトークン間で引き継ぐのは固定サイズの隠れ状態だけなので、コンテキスト長が伸びてもメモリ使用量は増加しません。

このリポジトリには、RWKV-7のリファレンス学習実装（train_temp）、高速なCUDAカーネル、GPT風・RNN・両者を組み合わせた各モードのデモスクリプト、旧バージョンのRWKV-4、5、6のコードが含まれています。また、jsonlからbinidx形式へのデータ準備、DeepSpeedとPyTorch Lightningを使った学習環境、安定した学習を再現するための初期化とweight decayのルールについても説明されています。

関連するエコシステムには、RWKVのpipパッケージ、GGUF形式の重み、モバイルおよびWebGPUでの推論プロジェクト、さらに他のコントリビューターが開発するPEFTやRLHFの実装などのfine-tuningツールがあります。

主な機能

  • 01GPT風の並列学習と、RNN型のメモリ使用量一定の推論に対応。
  • 02RWKV-7のリファレンス学習実装とCUDAカーネルを提供。
  • 03GPT、RNN、両者を組み合わせた推論モードのデモスクリプトを収録。
  • 04jsonlからbinidxへのデータ準備と、DeepSpeedを使った学習について説明。
  • 05ゼロから安定して学習するための初期化とweight decayの指針を提示。

ユースケース

  • Attentionを使わない、または時間計算量が線形の系列モデルを研究する。
  • リファレンスコードを使い、単一GPUで独自のRWKV言語モデルを学習する。
  • KV cacheのメモリ使用量が制約となる、長いコンテキストでの推論を実行する。
  • エッジやオンデバイスへのデプロイに向けて、Transformerに代わるRNNベースのモデルを検討する。

クイックスタート

  1. 1

    CUDA 12.1 の wheel を使って PyTorch をインストールする

    pip install torch --upgrade --extra-index-url https://download.pytorch.org/whl/cu121
  2. 2

    学習用の依存関係をインストールする（pytorch-lightning は 1.9.5 のままにする）

    pip install pytorch-lightning==1.9.5 deepspeed wandb ninja --upgrade
  3. 3

    RWKV-7 の参照学習ディレクトリに移動する

    cd RWKV-v7/train_temp/
  4. 4

    MiniPile データを train_temp/data にダウンロードし、準備スクリプトを実行する

    sh ./demo-training-prepare.sh

使う前に知っておきたいこと

  • 学習には NVIDIA CUDA GPU が必要です。デフォルト設定では約 7 GB の VRAM が必要です。
  • 参照学習コードでは pytorch-lightning が 1.9.5 に固定されており、DeepSpeed と CUDA カーネルに依存しています。
  • README によると、FLA の RWKV-7 実装は参照実装と整合しておらず、性能も劣ります。また、簡単に置き換えられる RWKV-7 レイヤーはありません。

トピック

#rnn#language-model#linear-attention#pytorch#llm#deep-learning#transformer-alternative#attention-mechanism#chatgpt#gpt

よくある質問

RWKV-LM とは？
GPTのTransformerのように並列学習できる、RNN型言語モデルアーキテクチャRWKVの公式学習・リファレンスコード。 RWKV-LMは、再帰型ニューラルネットワークでありながらTransformerのように並列学習できる言語モデルアーキテクチャRWKVの、メインとなる学習・リファレンスリポジトリです。現在のプロジェクトはRWKV-7「Goose」で、READMEでは、Attentionを使わず、時間計算量が線形で、メモリ使用量が一定、KV cacheも不要な設計として説明されています。推論時にトークン間で引き継ぐのは固定サイズの隠れ状態だけなので、コンテキスト長が伸びてもメモリ使用量は増加しません。
RWKV-LM は無料で使えますか？ライセンスは？
RWKV-LM は Apache-2.0 ライセンスのオープンソースです。商用製品で使う前にライセンス条件を確認してください。
RWKV-LM は今もメンテナンスされていますか？
活発にメンテナンス中。RWKV-LM の最終コミットは 2026年9月28日 です。
RWKV-LM は何に使えますか？
Attentionを使わない、または時間計算量が線形の系列モデルを研究する。 リファレンスコードを使い、単一GPUで独自のRWKV言語モデルを学習する。 KV cacheのメモリ使用量が制約となる、長いコンテキストでの推論を実行する。 エッジやオンデバイスへのデプロイに向けて、Transformerに代わるRNNベースのモデルを検討する。
RWKV-LM のインストール方法は？
CUDA 12.1 の wheel を使って PyTorch をインストールする: pip install torch --upgrade --extra-index-url https://download.pytorch.org/whl/cu121 学習用の依存関係をインストールする（pytorch-lightning は 1.9.5 のままにする）: pip install pytorch-lightning==1.9.5 deepspeed wandb ninja --upgrade RWKV-7 の参照学習ディレクトリに移動する: cd RWKV-v7/train_temp/ MiniPile データを train_temp/data にダウンロードし、準備スクリプトを実行する: sh ./demo-training-prepare.sh
RWKV-LM はどこで入手できますか？
ソースコードとインストール手順は GitHub にあります：https://github.com/BlinkDL/RWKV-LM

概要はプロジェクトの README をもとに AI が生成しています。GitHub の数値は毎日更新されます。

関連プロジェクト