ARWKV R1 7B

🚀 ARWKV🪿

ARWKV是一個基於純RNN的模型，結合了RWKV-7的時間混合機制和Transformer的MLP，具有高效的循環機制、無需自注意力機制、恆定的顯存使用和單GPU可訓練性等特點。本項目還提供了在AMD Radeon GPU上的推理方法，以及模型的使用示例和性能基準測試。

🚀 快速開始

環境準備

在AMD Radeon GPU上使用llama.cpp進行推理，可按以下步驟操作：

git clone https://github.com/MollySophia/llama.cpp.git -b rwkv-v7
cd llama.cpp
HIPCXX="$(hipconfig -l)/clang" HIP_PATH="$(hipconfig -R)" \
    cmake -S . -B build -DGGML_HIP=ON -DAMDGPU_TARGETS=gfx1030 -DCMAKE_BUILD_TYPE=Release \
    && cmake --build build --config Release -- -j 16
cd ./build/bin 

模型轉換

將safetensor模型轉換為gguf格式：

python ./convert_hf_to_gguf.py [model_dir]

模型量化

對模型進行量化：

./llama-quantize [model_dir] [Quantization accuracy]

在Webui中推理

使用llama-server在Webui中推理模型：

/llama-server -m [model_dir] -t [use_cpu_thread_number] -ngl 99 --host [host_number] --port [port_number]

注意：Radeon 7000系列使用gfx1100，Radeon 6000系列使用gfx1030。

在Nvidia GPU上推理

git clone https://github.com/MollySophia/llama.cpp.git -b rwkv-v7
cd llama.cpp
cmake -B build -DGGML_CUDA=ON
cmake --build build --config Release
cd ./build/bin 

模型轉換

python ./convert_hf_to_gguf.py [model_dir]

模型量化

./llama-quantize [model_dir] [Quantization_accuracy]

在Webui中推理

/llama-server -m [model_dir] -t [use_cpu_thread_number] -ngl 99 --host [host_number] --port [port_number]

✨ 主要特性

模型特性

• 高效循環機制：採用RWKV-7的高效循環機制。
• 無自注意力機制：完全O(n)複雜度，無需自注意力機制。
• 恆定顯存使用：顯存使用恆定，可在單GPU上訓練。

未來規劃

• 長上下文能力：即將開源支持16k+上下文的增強版本。
• 數學能力提升：針對數學問題進行特定改進。
• 強化推理模型：通過強化學習提升推理能力。

📦 安裝指南

pip3 install --upgrade transformers rwkv-fla

訓練前需設置環境變量：

export WKV_MODE=chunk

💻 使用示例

基礎用法

from transformers import AutoModelForCausalLM, AutoTokenizer
import threading

model = AutoModelForCausalLM.from_pretrained(
    "RWKV-Red-Team/ARWKV-R1-7B",
    device_map="auto",
    torch_dtype=torch.float16,
    trust_remote_code=True,
)
tokenizer = AutoTokenizer.from_pretrained(
    "RWKV-Red-Team/ARWKV-R1-7B"
)

system_prompt = "You are a world class trivia AI - provide accurate, succinct responses. "
messages = [
    {"role": "system", "content": system_prompt},
    {"role": "user", "content": prompt}]

text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
text = text + "<think>"
print(text)
model_inputs = tokenizer([text], return_tensors="pt").to(device)

streamer = TextIteratorStreamer(tokenizer, skip_prompt=False, skip_special_tokens=False)

generation_kwargs = dict(model_inputs, streamer=streamer, max_new_tokens=8192, do_sample=True,tokenizer=tokenizer,stop_strings=["<｜end▁of▁sentence｜>"])
thread = threading.Thread(target=model.generate, kwargs=generation_kwargs)
thread.start()

print("Streaming output:")
for new_text in streamer:
    print(new_text, end="", flush=True)

thread.join()

輸出示例

<｜begin▁of▁sentence｜>You are a world class trivia AI - provide accurate, succinct responses. <｜User｜>The world's largest rainforest, home to approximately three million species of plants and animals, is named after which river?<｜Assistant｜><think>
Okay, so I'm trying to solve this question about the world's largest rainforest and which river it's named after. Hmm, first, I think rainforest names often have links related to the region it's in. The most famous rainforest in the world is the Amazon. I remember hearing a lot about it being called that because rainforests are connected to specific river systems. 

Now, I'm trying to recall which river is named after the Amazon. I think it's the Amazon River. But I want to be sure. Let me see... the Amazon is a major rainforest located in South America. The Amazon River flows through it, which is why it's named after it. That makes sense because it's a very important river. I recall reading somewhere that all the rainforests are named after rivers related to their regions. So if the Amazon is named after its River, then the name would naturally be related to its source.

I wonder if it's the Amazon itself that's named after it, or another river named after it. But the official name for the Amazon is the Amazon Rainforest. The most significant rainforest in the world is the Amazon, and its name probably started with river-sounding names.
</think>

The largest rainforest located in South America is the Amazon. It is named after the river named after it, which is the Amazon River. Therefore, the Amazon River is the name given to the Amazon Rain Forest.

📚 詳細文檔

性能基準測試

關鍵特性

架構亮點

核心修改流程

Transformer Decoder Layer:
- Multi-head Latent Attention(MLA)
+ RWKV-7 Time Mixing (Eq.3)
- RoPE Positional Encoding
+ State Recurrence
= Hybrid Layer Output

📄 許可證

本項目採用Apache-2.0許可證。