OpenMOSE commited on
Commit
4adaf5e
·
verified ·
1 Parent(s): 98a8477

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +94 -3
README.md CHANGED
@@ -1,3 +1,94 @@
1
- ---
2
- license: apache-2.0
3
- ---
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: apache-2.0
3
+ ---
4
+
5
+ # RWKV-Reka-3.1-Flash
6
+
7
+
8
+
9
+ <div align="center">
10
+ <img src="./hxa079.png" style="border-radius: 15px; width: 60%; height: 60%; object-fit: cover; box-shadow: 10px 10px 20px rgba(0, 0, 0, 0.5); border: 2px solid white;" alt="PRWKV" />
11
+ </div>
12
+
13
+ ### Model Description
14
+
15
+ HRWKV7-Reka-3.1 Flash is an RNN hybrid architecture model that combines RWKV v7's linear attention mechanism with Group Query Attention (GQA) layers. Built upon the Reka-flash3.1 21B foundation, this model replaces most Transformer attention blocks with RWKV blocks while strategically maintaining some GQA layers to enhance performance on specific tasks.
16
+
17
+ - **Developed by:** OpenMOSE
18
+ - **Model type:** Hybrid Linear-Attention Language Model
19
+ - **Language(s):** Multilingual (inherited from Reka-flash3.1 21B)
20
+ - **License:** Apache-2.0
21
+ - **Base Model:** Reka-flash3 21B(https://huggingface.co/RekaAI/reka-flash-3.1)
22
+ - **Year:** 2025
23
+
24
+ ### Architecture Specifications
25
+
26
+ - **Architecture:** RWKV v7 based "hxa079" Architecture + Group Query Attention Hybrid
27
+ - **Total Layers:** 44 layers (L44D6114)
28
+ - 38 RWKV layers (with Rope)
29
+ - 6 GQA layers (No Rope, No Position Embeddings)
30
+ - **Hidden Dimension:** 6144
31
+ - **Training Context Window:** 8192 tokens
32
+ - **Inference Context Window** 40000+
33
+ - **Training Strategy** Following RADLADS method based knowledge distillation
34
+
35
+ ## Technical Innovation
36
+
37
+ ### RWKV "hxa079" Architecture
38
+
39
+ The model implements several key improvements over original RWKV architectures:
40
+
41
+ 1. **Token Shift Removal**: In order to effectively inherit the teacher model weights, we removed the residual connection one token ago.
42
+ 2. **GroupNorm Removal**: Helps improve training stability issues
43
+ 3. **k_first Introduction**: Experimentally adopted the approach of residually connecting k layers in layer 0.
44
+
45
+ ### Hybrid Design Benefits
46
+
47
+ - **Linear Attention Inference**: RWKV blocks enable O(1) memory complexity during inference, and the hybrid approach reduces the KVCache to 1/7 of full GQA.
48
+ - **Enhanced Needle Tasks**: Strategic placement of GQA layers significantly improves performance on needle-in-haystack retrieval tasks, addressing a known limitation of pure linear attention models
49
+ - **Implicit Position Encoding**: Interestingly, the model achieves better performance when RoPE (Rotary Position Embedding) is not applied to GQA layers, suggesting that RWKV blocks provide implicit positional encoding capabilities
50
+
51
+ ## Intended Use
52
+
53
+ This is an **experimental research model** designed to explore hybrid architectures combining linear and quadratic attention mechanisms. It is intended for:
54
+
55
+ - Research into efficient attention mechanisms
56
+ - Benchmarking hybrid architecture performance
57
+ - Exploring linear attention limitations and solutions
58
+ - Academic and industrial R&D purposes
59
+
60
+ ## Limitations
61
+
62
+ - **Experimental Status**: This model is in experimental stages and may exhibit unexpected behaviors
63
+ - **Context Window**: Limited to 8192 tokens during training, though RWKV architecture theoretically supports longer sequences
64
+ - **Performance Variability**: As a hybrid model, performance may vary significantly across different task types
65
+
66
+ ## Training Details
67
+
68
+ - **Training Context Window:** 8192 tokens
69
+ - **Training GPU** AMD Instinct MI300X (takes 290hrs) AMD Developer Cloud.(Thank you for credit support)
70
+ - **Training Strategy** 8bit MLP Quant, frozen emb,mlp,head, Deepspeed Stage1, Stage1 100M, Stage2(ctx4096) 360M Stage3(ctx8192) 300M
71
+ - **Training Stage** Stage3 - Reduced temperature stepped knowledge distillation.(stage3 final temp=0.7)
72
+ - **Base Model Initialization:** Weights initialized from Reka-3.1-flash 21B
73
+ - **Architecture Conversion:** Transformer attention blocks systematically replaced with RWKV blocks, except for 6 strategically placed GQA layers
74
+
75
+ ## Evaluation
76
+
77
+ Performance evaluation is ongoing. The model shows promising results in:
78
+ - Maintaining base model capabilities while achieving linear attention efficiency
79
+ - Significantly improved needle-in-haystack task performance compared to pure RWKV architectures
80
+ - Competitive performance on standard language modeling benchmarks
81
+
82
+ ## Usage with Hugging Face Transformers
83
+
84
+ urrently in development. stay tuned :)
85
+
86
+ ## Code Repositories
87
+
88
+ - **RADLADS Project Code:** The main codebase for the RADLADS paper, including conversion scripts and model code, can be found at: [https://github.com/recursal/RADLADS](https://github.com/recursal/RADLADS)
89
+ - **ARWKV Project Code** The ARWKV original training code, can be found at: [https://github.com/yynil/RWKVInside](https://github.com/yynil/RWKVInside)
90
+ - **Specific Training Code (OpenMOSE):** The training code for this particular `RWKV-Reka-3.1-Flash` model is available at: [https://github.com/OpenMOSE/RWKVInside](https://github.com/OpenMOSE/RWKVInside) (Note: this repository is still under development and may contain bugs.)
91
+
92
+ ## Model Card Contact
93
+
94
+ OpenMOSE - 2025