RWKV Architecture Integration into Hugging Face Transformers

TL;DR

RWKV, a novel RNN‑based architecture that mimics transformer attention, is now officially supported in the Hugging Face transformers library, giving developers access to open‑source language models that handle very long contexts with the speed and memory efficiency of RNNs.

Overview of the RWKV Project

The RWKV project is led by Bo Peng (GitHub: BlinkDL) and is maintained by an active Discord community. Stability AI donated the GPUs used for training. The project’s roadmap includes performance improvements (e.g., RWKV.cpp, quantization), scalability enhancements (dataset processing), and research extensions such as chat‑fine‑tuning and multimodal finetuning. Community members can join the official Discord channel to contribute.

How RWKV Bridges RNNs and Transformers

RNN limitations and transformer advantages

  • Traditional RNNs reuse the same weights at each time step, which leads to vanishing‑gradient problems and poor long‑range memory. LSTMs and GRUs mitigate this partially but still struggle with very long sequences.
  • Transformers process all tokens in parallel via self‑attention, using query, key, and value projections to compute attention scores. This design resolves the long‑range dependency issue and speeds up training compared to classic RNNs.

RWKV’s hybrid design

  • RWKV is inspired by Apple’s Attention‑Free Transformer and has been simplified into an RNN‑compatible form.
  • It retains the transformer‑style embedding, layer‑norm, and causal language‑model head, but replaces the attention layer with a recurrence‑based formulation that yields the same expressive power as self‑attention.
  • Additional tricks such as TokenShift and SmallInitEmb (documented in the official GitHub README) are required for the model to match GPT‑level performance.

Technical Highlights of the RWKV Architecture

Long‑context capability

  • RWKV can handle context windows of 8 192 tokens (ctx8192) with the same inference speed and RAM usage as a 1 024‑token model.
  • Empirical loss curves show that larger context lengths improve language‑model loss across model sizes, demonstrating effective long‑range memorization.

Training efficiency

  • Unlike classic RNNs, RWKV can be trained in a “linearized GPT” fashion, allowing parallelism across batches and faster convergence than traditional recurrent models.
  • The current training pipeline scales up to 14 B parameters, with ongoing fixes for numerical stability in the RWKV‑4 series.

Available Model Checkpoints

Pure language models (RWKV‑4)

  • Model sizes range from ~170 M to 14 B parameters.
  • All models are pretrained on The Pile dataset and have been benchmarked against state‑of‑the‑art baselines, showing comparable performance.

Instruction‑fine‑tuned chat models (RWKV‑4 Raven)

  • The Raven series fine‑tunes RWKV‑4 on instruction datasets such as ALPACA, CodeAlpaca, Guanaco, GPT‑4All, and ShareGPT.
  • Variants exist for different language mixes (English only, English + Chinese + Japanese, etc.) and sizes (1.5 B, 7 B, 14 B).
  • All checkpoints are hosted under the RWKV organization on the Hugging Face Hub.

Using RWKV with 🤗 Transformers

Text generation example

from transformers import pipeline
model_id = "RWKV/rwkv-4-169m-pile"
pipe = pipeline("text-generation", model=model_id)
print(pipe("In a shocking finding, scientist discovered a herd of dragons...", max_new_tokens=20))

The pipeline returns coherent continuations comparable to transformer‑based generators.

Chat model (Raven) example

from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "RWKV/rwkv-raven-1b5"
model = AutoModelForCausalLM.from_pretrained(model_id).to(0)
tokenizer = AutoTokenizer.from_pretrained(model_id)
prompt = "### Instruction: Tell me about ravens\n### Response:"
inputs = tokenizer(prompt, return_tensors="pt").to(0)
output = model.generate(inputs["input_ids"], max_new_tokens=100)
print(tokenizer.decode(output[0], skip_special_tokens=True))

The model follows the Alpaca‑style instruction format and produces detailed responses.

Converting Original RWKV Weights to Hugging Face Format

A conversion script (convert_rwkv_checkpoint_to_hf.py) is bundled with the transformers repo. Users upload the raw checkpoint to a Hub repo, then run:

python convert_rwkv_checkpoint_to_hf.py \
  --repo_id RAW_HUB_REPO \
  --checkpoint_file RAW_FILE \
  --output_dir OUTPUT_DIR

Adding --push_to_hub and --model_name uploads the converted model directly to the Hub.

Future Directions

  • Multilingual RWKV – Work is underway on a multilingual corpus and tokenizer, expanding the model’s language coverage.
  • Community research – The Discord channel hosts projects on novel training recipes, benchmarking, and architectural tweaks.
  • Compression & acceleration – Because RWKV relies only on matrix‑vector operations, it is well‑suited for quantization (4‑bit/8‑bit), ONNX export, and experimental hardware such as photonic accelerators. Integration with the optimum library and repositories like rwkv.cpp and rwkv-cpp-cuda will further speed up inference.

Acknowledgements

The Hugging Face team thanks Bo Peng, the RWKV community, and contributors such as Johan Wind (RWKV blog posts), ArEnSc (initial Transformers PR), Merve Noyan, Maria Khalusova, and Pedro Cuenca for reviewing and supporting the integration.

Citation

If you use RWKV in research, cite the project using the provided CITATION.cff file in the RWKV‑LM repository.

Sources