Instructions to use aabbdev/RWKV7-1.5B-20260805 with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- Transformers
How to use aabbdev/RWKV7-1.5B-20260805 with Transformers:
# Use a pipeline as a high-level helper from transformers import pipeline pipe = pipeline("text-generation", model="aabbdev/RWKV7-1.5B-20260805", trust_remote_code=True) messages = [ {"role": "user", "content": "Who are you?"}, ] pipe(messages)# Load model directly from transformers import AutoModelForCausalLM model = AutoModelForCausalLM.from_pretrained("aabbdev/RWKV7-1.5B-20260805", trust_remote_code=True, device_map="auto") - RWKV
How to use aabbdev/RWKV7-1.5B-20260805 with RWKV:
# No code snippets available yet for this library. # To use this model, check the repository files and the library's documentation. # Want to help? PRs adding snippets are welcome at: # https://github.com/huggingface/huggingface.js
- Notebooks
- Google Colab
- Kaggle
- Local Apps Settings
- vLLM
How to use aabbdev/RWKV7-1.5B-20260805 with vLLM:
Install from pip and serve model
# Install vLLM from pip: pip install vllm # Start the vLLM server: vllm serve "aabbdev/RWKV7-1.5B-20260805" # Call the server using curl (OpenAI-compatible API): curl -X POST "http://localhost:8000/v1/chat/completions" \ -H "Content-Type: application/json" \ --data '{ "model": "aabbdev/RWKV7-1.5B-20260805", "messages": [ { "role": "user", "content": "What is the capital of France?" } ] }'Use Docker
docker model run hf.co/aabbdev/RWKV7-1.5B-20260805
- SGLang
How to use aabbdev/RWKV7-1.5B-20260805 with SGLang:
Install from pip and serve model
# Install SGLang from pip: pip install sglang # Start the SGLang server: python3 -m sglang.launch_server \ --model-path "aabbdev/RWKV7-1.5B-20260805" \ --host 0.0.0.0 \ --port 30000 # Call the server using curl (OpenAI-compatible API): curl -X POST "http://localhost:30000/v1/chat/completions" \ -H "Content-Type: application/json" \ --data '{ "model": "aabbdev/RWKV7-1.5B-20260805", "messages": [ { "role": "user", "content": "What is the capital of France?" } ] }'Use Docker images
docker run --gpus all \ --shm-size 32g \ -p 30000:30000 \ -v ~/.cache/huggingface:/root/.cache/huggingface \ --env "HF_TOKEN=<secret>" \ --ipc=host \ lmsysorg/sglang:latest \ python3 -m sglang.launch_server \ --model-path "aabbdev/RWKV7-1.5B-20260805" \ --host 0.0.0.0 \ --port 30000 # Call the server using curl (OpenAI-compatible API): curl -X POST "http://localhost:30000/v1/chat/completions" \ -H "Content-Type: application/json" \ --data '{ "model": "aabbdev/RWKV7-1.5B-20260805", "messages": [ { "role": "user", "content": "What is the capital of France?" } ] }' - Docker Model Runner
How to use aabbdev/RWKV7-1.5B-20260805 with Docker Model Runner:
docker model run hf.co/aabbdev/RWKV7-1.5B-20260805
Publish RWKV7-1.5B-20260805
Browse files- .gitattributes +2 -34
- LICENSE +201 -0
- NOTICE +4 -0
- README.md +279 -0
- chat_template.jinja +56 -0
- config.json +31 -0
- configuration_rwkv7.py +153 -0
- generation_config.json +6 -0
- inference/generate.py +239 -0
- inference/kernel.py +0 -0
- inference/model_loader.py +178 -0
- inference/requirements.txt +5 -0
- inference/runtime.py +0 -0
- model.safetensors +3 -0
- modeling_rwkv7.py +1186 -0
- release-manifest.json +332 -0
- tokenizer.json +0 -0
- tokenizer_config.json +18 -0
.gitattributes
CHANGED
|
@@ -1,35 +1,3 @@
|
|
| 1 |
-
*.7z filter=lfs diff=lfs merge=lfs -text
|
| 2 |
-
*.arrow filter=lfs diff=lfs merge=lfs -text
|
| 3 |
-
*.bin filter=lfs diff=lfs merge=lfs -text
|
| 4 |
-
*.bz2 filter=lfs diff=lfs merge=lfs -text
|
| 5 |
-
*.ckpt filter=lfs diff=lfs merge=lfs -text
|
| 6 |
-
*.ftz filter=lfs diff=lfs merge=lfs -text
|
| 7 |
-
*.gz filter=lfs diff=lfs merge=lfs -text
|
| 8 |
-
*.h5 filter=lfs diff=lfs merge=lfs -text
|
| 9 |
-
*.joblib filter=lfs diff=lfs merge=lfs -text
|
| 10 |
-
*.lfs.* filter=lfs diff=lfs merge=lfs -text
|
| 11 |
-
*.mlmodel filter=lfs diff=lfs merge=lfs -text
|
| 12 |
-
*.model filter=lfs diff=lfs merge=lfs -text
|
| 13 |
-
*.msgpack filter=lfs diff=lfs merge=lfs -text
|
| 14 |
-
*.npy filter=lfs diff=lfs merge=lfs -text
|
| 15 |
-
*.npz filter=lfs diff=lfs merge=lfs -text
|
| 16 |
-
*.onnx filter=lfs diff=lfs merge=lfs -text
|
| 17 |
-
*.ot filter=lfs diff=lfs merge=lfs -text
|
| 18 |
-
*.parquet filter=lfs diff=lfs merge=lfs -text
|
| 19 |
-
*.pb filter=lfs diff=lfs merge=lfs -text
|
| 20 |
-
*.pickle filter=lfs diff=lfs merge=lfs -text
|
| 21 |
-
*.pkl filter=lfs diff=lfs merge=lfs -text
|
| 22 |
-
*.pt filter=lfs diff=lfs merge=lfs -text
|
| 23 |
-
*.pth filter=lfs diff=lfs merge=lfs -text
|
| 24 |
-
*.rar filter=lfs diff=lfs merge=lfs -text
|
| 25 |
*.safetensors filter=lfs diff=lfs merge=lfs -text
|
| 26 |
-
|
| 27 |
-
*.
|
| 28 |
-
*.tar filter=lfs diff=lfs merge=lfs -text
|
| 29 |
-
*.tflite filter=lfs diff=lfs merge=lfs -text
|
| 30 |
-
*.tgz filter=lfs diff=lfs merge=lfs -text
|
| 31 |
-
*.wasm filter=lfs diff=lfs merge=lfs -text
|
| 32 |
-
*.xz filter=lfs diff=lfs merge=lfs -text
|
| 33 |
-
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
-
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
-
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
*.safetensors filter=lfs diff=lfs merge=lfs -text
|
| 2 |
+
*.pth filter=lfs diff=lfs merge=lfs -text
|
| 3 |
+
*.bin filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
LICENSE
ADDED
|
@@ -0,0 +1,201 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
Apache License
|
| 2 |
+
Version 2.0, January 2004
|
| 3 |
+
http://www.apache.org/licenses/
|
| 4 |
+
|
| 5 |
+
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
| 6 |
+
|
| 7 |
+
1. Definitions.
|
| 8 |
+
|
| 9 |
+
"License" shall mean the terms and conditions for use, reproduction,
|
| 10 |
+
and distribution as defined by Sections 1 through 9 of this document.
|
| 11 |
+
|
| 12 |
+
"Licensor" shall mean the copyright owner or entity authorized by
|
| 13 |
+
the copyright owner that is granting the License.
|
| 14 |
+
|
| 15 |
+
"Legal Entity" shall mean the union of the acting entity and all
|
| 16 |
+
other entities that control, are controlled by, or are under common
|
| 17 |
+
control with that entity. For the purposes of this definition,
|
| 18 |
+
"control" means (i) the power, direct or indirect, to cause the
|
| 19 |
+
direction or management of such entity, whether by contract or
|
| 20 |
+
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
| 21 |
+
outstanding shares, or (iii) beneficial ownership of such entity.
|
| 22 |
+
|
| 23 |
+
"You" (or "Your") shall mean an individual or Legal Entity
|
| 24 |
+
exercising permissions granted by this License.
|
| 25 |
+
|
| 26 |
+
"Source" form shall mean the preferred form for making modifications,
|
| 27 |
+
including but not limited to software source code, documentation
|
| 28 |
+
source, and configuration files.
|
| 29 |
+
|
| 30 |
+
"Object" form shall mean any form resulting from mechanical
|
| 31 |
+
transformation or translation of a Source form, including but
|
| 32 |
+
not limited to compiled object code, generated documentation,
|
| 33 |
+
and conversions to other media types.
|
| 34 |
+
|
| 35 |
+
"Work" shall mean the work of authorship, whether in Source or
|
| 36 |
+
Object form, made available under the License, as indicated by a
|
| 37 |
+
copyright notice that is included in or attached to the work
|
| 38 |
+
(an example is provided in the Appendix below).
|
| 39 |
+
|
| 40 |
+
"Derivative Works" shall mean any work, whether in Source or Object
|
| 41 |
+
form, that is based on (or derived from) the Work and for which the
|
| 42 |
+
editorial revisions, annotations, elaborations, or other modifications
|
| 43 |
+
represent, as a whole, an original work of authorship. For the purposes
|
| 44 |
+
of this License, Derivative Works shall not include works that remain
|
| 45 |
+
separable from, or merely link (or bind by name) to the interfaces of,
|
| 46 |
+
the Work and Derivative Works thereof.
|
| 47 |
+
|
| 48 |
+
"Contribution" shall mean any work of authorship, including
|
| 49 |
+
the original version of the Work and any modifications or additions
|
| 50 |
+
to that Work or Derivative Works thereof, that is intentionally
|
| 51 |
+
submitted to Licensor for inclusion in the Work by the copyright owner
|
| 52 |
+
or by an individual or Legal Entity authorized to submit on behalf of
|
| 53 |
+
the copyright owner. For the purposes of this definition, "submitted"
|
| 54 |
+
means any form of electronic, verbal, or written communication sent
|
| 55 |
+
to the Licensor or its representatives, including but not limited to
|
| 56 |
+
communication on electronic mailing lists, source code control systems,
|
| 57 |
+
and issue tracking systems that are managed by, or on behalf of, the
|
| 58 |
+
Licensor for the purpose of discussing and improving the Work, but
|
| 59 |
+
excluding communication that is conspicuously marked or otherwise
|
| 60 |
+
designated in writing by the copyright owner as "Not a Contribution."
|
| 61 |
+
|
| 62 |
+
"Contributor" shall mean Licensor and any individual or Legal Entity
|
| 63 |
+
on behalf of whom a Contribution has been received by Licensor and
|
| 64 |
+
subsequently incorporated within the Work.
|
| 65 |
+
|
| 66 |
+
2. Grant of Copyright License. Subject to the terms and conditions of
|
| 67 |
+
this License, each Contributor hereby grants to You a perpetual,
|
| 68 |
+
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
| 69 |
+
copyright license to reproduce, prepare Derivative Works of,
|
| 70 |
+
publicly display, publicly perform, sublicense, and distribute the
|
| 71 |
+
Work and such Derivative Works in Source or Object form.
|
| 72 |
+
|
| 73 |
+
3. Grant of Patent License. Subject to the terms and conditions of
|
| 74 |
+
this License, each Contributor hereby grants to You a perpetual,
|
| 75 |
+
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
| 76 |
+
(except as stated in this section) patent license to make, have made,
|
| 77 |
+
use, offer to sell, sell, import, and otherwise transfer the Work,
|
| 78 |
+
where such license applies only to those patent claims licensable
|
| 79 |
+
by such Contributor that are necessarily infringed by their
|
| 80 |
+
Contribution(s) alone or by combination of their Contribution(s)
|
| 81 |
+
with the Work to which such Contribution(s) was submitted. If You
|
| 82 |
+
institute patent litigation against any entity (including a
|
| 83 |
+
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
| 84 |
+
or a Contribution incorporated within the Work constitutes direct
|
| 85 |
+
or contributory patent infringement, then any patent licenses
|
| 86 |
+
granted to You under this License for that Work shall terminate
|
| 87 |
+
as of the date such litigation is filed.
|
| 88 |
+
|
| 89 |
+
4. Redistribution. You may reproduce and distribute copies of the
|
| 90 |
+
Work or Derivative Works thereof in any medium, with or without
|
| 91 |
+
modifications, and in Source or Object form, provided that You
|
| 92 |
+
meet the following conditions:
|
| 93 |
+
|
| 94 |
+
(a) You must give any other recipients of the Work or
|
| 95 |
+
Derivative Works a copy of this License; and
|
| 96 |
+
|
| 97 |
+
(b) You must cause any modified files to carry prominent notices
|
| 98 |
+
stating that You changed the files; and
|
| 99 |
+
|
| 100 |
+
(c) You must retain, in the Source form of any Derivative Works
|
| 101 |
+
that You distribute, all copyright, patent, trademark, and
|
| 102 |
+
attribution notices from the Source form of the Work,
|
| 103 |
+
excluding those notices that do not pertain to any part of
|
| 104 |
+
the Derivative Works; and
|
| 105 |
+
|
| 106 |
+
(d) If the Work includes a "NOTICE" text file as part of its
|
| 107 |
+
distribution, then any Derivative Works that You distribute must
|
| 108 |
+
include a readable copy of the attribution notices contained
|
| 109 |
+
within such NOTICE file, excluding those notices that do not
|
| 110 |
+
pertain to any part of the Derivative Works, in at least one
|
| 111 |
+
of the following places: within a NOTICE text file distributed
|
| 112 |
+
as part of the Derivative Works; within the Source form or
|
| 113 |
+
documentation, if provided along with the Derivative Works; or,
|
| 114 |
+
within a display generated by the Derivative Works, if and
|
| 115 |
+
wherever such third-party notices normally appear. The contents
|
| 116 |
+
of the NOTICE file are for informational purposes only and do not
|
| 117 |
+
modify the License. You may add Your own attribution notices
|
| 118 |
+
within Derivative Works that You distribute, alongside or as an
|
| 119 |
+
addendum to the NOTICE text from the Work, provided that such
|
| 120 |
+
additional attribution notices cannot be construed as modifying
|
| 121 |
+
the License.
|
| 122 |
+
|
| 123 |
+
You may add Your own copyright statement to Your modifications and
|
| 124 |
+
may provide additional or different license terms and conditions
|
| 125 |
+
for use, reproduction, or distribution of Your modifications, or
|
| 126 |
+
for any such Derivative Works as a whole, provided Your use,
|
| 127 |
+
reproduction, and distribution of the Work otherwise complies with
|
| 128 |
+
the conditions stated in this License.
|
| 129 |
+
|
| 130 |
+
5. Submission of Contributions. Unless You explicitly state otherwise,
|
| 131 |
+
any Contribution intentionally submitted for inclusion in the Work
|
| 132 |
+
by You to the Licensor shall be under the terms and conditions of
|
| 133 |
+
this License, without any additional terms or conditions.
|
| 134 |
+
Notwithstanding the above, nothing herein shall supersede or modify
|
| 135 |
+
the terms of any separate license agreement you may have executed
|
| 136 |
+
with Licensor regarding such Contributions.
|
| 137 |
+
|
| 138 |
+
6. Trademarks. This License does not grant permission to use the trade
|
| 139 |
+
names, trademarks, service marks, or product names of the Licensor,
|
| 140 |
+
except as required for reasonable and customary use in describing the
|
| 141 |
+
origin of the Work and reproducing the content of the NOTICE file.
|
| 142 |
+
|
| 143 |
+
7. Disclaimer of Warranty. Unless required by applicable law or
|
| 144 |
+
agreed to in writing, Licensor provides the Work (and each
|
| 145 |
+
Contributor provides its Contributions) on an "AS IS" BASIS,
|
| 146 |
+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
| 147 |
+
implied, including, without limitation, any warranties or conditions
|
| 148 |
+
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
| 149 |
+
PARTICULAR PURPOSE. You are solely responsible for determining the
|
| 150 |
+
appropriateness of using or redistributing the Work and assume any
|
| 151 |
+
risks associated with Your exercise of permissions under this License.
|
| 152 |
+
|
| 153 |
+
8. Limitation of Liability. In no event and under no legal theory,
|
| 154 |
+
whether in tort (including negligence), contract, or otherwise,
|
| 155 |
+
unless required by applicable law (such as deliberate and grossly
|
| 156 |
+
negligent acts) or agreed to in writing, shall any Contributor be
|
| 157 |
+
liable to You for damages, including any direct, indirect, special,
|
| 158 |
+
incidental, or consequential damages of any character arising as a
|
| 159 |
+
result of this License or out of the use or inability to use the
|
| 160 |
+
Work (including but not limited to damages for loss of goodwill,
|
| 161 |
+
work stoppage, computer failure or malfunction, or any and all
|
| 162 |
+
other commercial damages or losses), even if such Contributor
|
| 163 |
+
has been advised of the possibility of such damages.
|
| 164 |
+
|
| 165 |
+
9. Accepting Warranty or Additional Liability. While redistributing
|
| 166 |
+
the Work or Derivative Works thereof, You may choose to offer, and
|
| 167 |
+
charge a fee for, acceptance of support, warranty, indemnity, or
|
| 168 |
+
other liability obligations and/or rights consistent with this
|
| 169 |
+
License. However, in accepting such obligations, You may act only
|
| 170 |
+
on Your own behalf and on Your sole responsibility, not on behalf
|
| 171 |
+
of any other Contributor, and only if You agree to indemnify,
|
| 172 |
+
defend, and hold each Contributor harmless for any liability
|
| 173 |
+
incurred by, or claims asserted against, such Contributor by reason
|
| 174 |
+
of your accepting any such warranty or additional liability.
|
| 175 |
+
|
| 176 |
+
END OF TERMS AND CONDITIONS
|
| 177 |
+
|
| 178 |
+
APPENDIX: How to apply the Apache License to your work.
|
| 179 |
+
|
| 180 |
+
To apply the Apache License to your work, attach the following
|
| 181 |
+
boilerplate notice, with the fields enclosed by brackets "[]"
|
| 182 |
+
replaced with your own identifying information. (Don't include
|
| 183 |
+
the brackets!) The text should be enclosed in the appropriate
|
| 184 |
+
comment syntax for the file format. We also recommend that a
|
| 185 |
+
file or class name and description of purpose be included on the
|
| 186 |
+
same "printed page" as the copyright notice for easier
|
| 187 |
+
identification within third-party archives.
|
| 188 |
+
|
| 189 |
+
Copyright [yyyy] [name of copyright owner]
|
| 190 |
+
|
| 191 |
+
Licensed under the Apache License, Version 2.0 (the "License");
|
| 192 |
+
you may not use this file except in compliance with the License.
|
| 193 |
+
You may obtain a copy of the License at
|
| 194 |
+
|
| 195 |
+
http://www.apache.org/licenses/LICENSE-2.0
|
| 196 |
+
|
| 197 |
+
Unless required by applicable law or agreed to in writing, software
|
| 198 |
+
distributed under the License is distributed on an "AS IS" BASIS,
|
| 199 |
+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 200 |
+
See the License for the specific language governing permissions and
|
| 201 |
+
limitations under the License.
|
NOTICE
ADDED
|
@@ -0,0 +1,4 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
RWKV-7 model release
|
| 2 |
+
Source: BlinkDL/rwkv7-g1/rwkv7-g1i-1.5b-20260805-ctx16384.pth
|
| 3 |
+
Bundled Transformers RWKV-7 code: huggingface/transformers@4ad9ed0747ed6ba75c787e8f9040dcd64b166ee2 (Apache-2.0).
|
| 4 |
+
Exported inference bundle licensed under Apache-2.0.
|
README.md
ADDED
|
@@ -0,0 +1,279 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
library_name: transformers
|
| 3 |
+
pipeline_tag: text-generation
|
| 4 |
+
license: apache-2.0
|
| 5 |
+
language:
|
| 6 |
+
- "en"
|
| 7 |
+
- "zh"
|
| 8 |
+
- "fr"
|
| 9 |
+
- "es"
|
| 10 |
+
- "de"
|
| 11 |
+
- "pt"
|
| 12 |
+
- "ru"
|
| 13 |
+
- "it"
|
| 14 |
+
- "ja"
|
| 15 |
+
- "ko"
|
| 16 |
+
- "vi"
|
| 17 |
+
- "ar"
|
| 18 |
+
datasets:
|
| 19 |
+
- "HuggingFaceFW/fineweb-edu"
|
| 20 |
+
- "mlfoundations/dclm-baseline-1.0"
|
| 21 |
+
- "cerebras/SlimPajama-627B"
|
| 22 |
+
- "EleutherAI/pile"
|
| 23 |
+
- "bigcode/starcoderdata"
|
| 24 |
+
- "oscar-corpus/OSCAR-2301"
|
| 25 |
+
tags:
|
| 26 |
+
- rwkv
|
| 27 |
+
- rwkv7
|
| 28 |
+
- recurrent
|
| 29 |
+
- causal-lm
|
| 30 |
+
- conversational
|
| 31 |
+
---
|
| 32 |
+
|
| 33 |
+
<!-- markdownlint-disable first-line-h1 -->
|
| 34 |
+
<!-- markdownlint-disable html -->
|
| 35 |
+
|
| 36 |
+
<div align="center">
|
| 37 |
+
<a href="https://www.rwkv.com/">
|
| 38 |
+
<img src="https://www.rwkv.com/images/avatar.png" width="140" alt="RWKV logo" />
|
| 39 |
+
</a>
|
| 40 |
+
<h1>RWKV7-1.5B-20260805</h1>
|
| 41 |
+
<p><strong>RWKV-7 “Goose” · constant-state recurrent language modeling</strong></p>
|
| 42 |
+
</div>
|
| 43 |
+
|
| 44 |
+
<div align="center">
|
| 45 |
+
<a href="https://www.rwkv.com/"><img alt="Website" src="https://img.shields.io/badge/Website-RWKV-16a7c9" /></a>
|
| 46 |
+
<a href="https://huggingface.co/BlinkDL"><img alt="Hugging Face" src="https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-BlinkDL-ffc107" /></a>
|
| 47 |
+
<a href="https://github.com/BlinkDL/RWKV-LM"><img alt="GitHub" src="https://img.shields.io/badge/GitHub-RWKV--LM-181717?logo=github" /></a>
|
| 48 |
+
<a href="https://arxiv.org/abs/2503.14456v2"><img alt="RWKV-7 paper" src="https://img.shields.io/badge/Paper-arXiv%3A2503.14456-b31b1b" /></a>
|
| 49 |
+
<a href="LICENSE"><img alt="License" src="https://img.shields.io/badge/License-apache-2.0-4c8bf5" /></a>
|
| 50 |
+
</div>
|
| 51 |
+
|
| 52 |
+
---
|
| 53 |
+
|
| 54 |
+
## Model introduction
|
| 55 |
+
|
| 56 |
+
This is an official BlinkDL release of **RWKV-7 Goose** in Hugging Face
|
| 57 |
+
Transformers format. RWKV-7 is an attention-free recurrent architecture with a
|
| 58 |
+
constant-size recurrent state and constant inference work per generated token.
|
| 59 |
+
Training remains parallelizable.
|
| 60 |
+
|
| 61 |
+
This checkpoint is a **base model** pretrained with web, code, synthetic, instruction, chat, and reasoning data. It is suitable for evaluation, post-training, and fine-tuning; the included chat template is a prompt interface, not a claim that the checkpoint is a safety-aligned assistant.
|
| 62 |
+
|
| 63 |
+
The Transformers integration, conversion, release packaging, Fast Tokenizer, and
|
| 64 |
+
optional TileLang inference implementation are distributed with this release.
|
| 65 |
+
|
| 66 |
+
## Highlights
|
| 67 |
+
|
| 68 |
+
- **Constant recurrent state:** memory does not grow like an attention KV cache.
|
| 69 |
+
- **Bundled Transformers integration:** auditable remote configuration and modeling
|
| 70 |
+
modules provide generation, recurrent cache continuation, training, and LoRA
|
| 71 |
+
workflows on Transformers 5.15+.
|
| 72 |
+
- **Exact Fast Tokenizer:** self-contained Rust-backed `tokenizer.json`, generated
|
| 73 |
+
from the canonical RWKV World byte vocabulary during conversion.
|
| 74 |
+
- **Chat-ready:** `chat_template.jinja` supports system, multi-turn, thinking, and
|
| 75 |
+
strict model-generated tool-call prompts.
|
| 76 |
+
- **Optional optimized runtime:** the isolated [`inference/`](inference/) bundle
|
| 77 |
+
provides PyTorch fallback and TileLang acceleration without changing the
|
| 78 |
+
standard model root.
|
| 79 |
+
|
| 80 |
+
## Model overview
|
| 81 |
+
|
| 82 |
+
| Field | Value |
|
| 83 |
+
| --- | --- |
|
| 84 |
+
| Repository | `aabbdev/RWKV7-1.5B-20260805` |
|
| 85 |
+
| Architecture class | `Rwkv7ForCausalLM` |
|
| 86 |
+
| Public size label | `1.5`B |
|
| 87 |
+
| Source parameters | `1,527,668,736` |
|
| 88 |
+
| Serialized parameters | `1,527,668,736` |
|
| 89 |
+
| Synthesized compatibility tensors | `0` |
|
| 90 |
+
| Layers | `24` |
|
| 91 |
+
| Hidden / FFN size | `2048` / `8192` |
|
| 92 |
+
| Heads / head size | `32` / `64` |
|
| 93 |
+
| Vocabulary | `65536` |
|
| 94 |
+
| Training context | `16384 tokens` |
|
| 95 |
+
| Weight dtype | `bfloat16` |
|
| 96 |
+
| Numerical conversion | `source dtype preserved` |
|
| 97 |
+
| Metadata profile | `g1i` |
|
| 98 |
+
| Metadata provenance | `locked-profile` |
|
| 99 |
+
| Source checkpoint | [`BlinkDL/rwkv7-g1/rwkv7-g1i-1.5b-20260805-ctx16384.pth`](https://huggingface.co/BlinkDL/rwkv7-g1/blob/ede85bf8ab2e59aff7d7ca909fbbc73317866d89/rwkv7-g1i-1.5b-20260805-ctx16384.pth) |
|
| 100 |
+
| Source SHA-256 | `32ef7b5bf4dc8bde843cf26dfad809a1f527e2e76a9e790e7d406e71bcd785da` |
|
| 101 |
+
|
| 102 |
+
## Transformers quickstart
|
| 103 |
+
|
| 104 |
+
The repository includes `configuration_rwkv7.py` and `modeling_rwkv7.py`, adapted
|
| 105 |
+
from the Transformers RWKV-7 integration at commit
|
| 106 |
+
[`4ad9ed0`](https://github.com/huggingface/transformers/commit/4ad9ed0747ed6ba75c787e8f9040dcd64b166ee2).
|
| 107 |
+
Review those files and pin a model-repository revision in production. Passing
|
| 108 |
+
`trust_remote_code=True` selects this bundled implementation even when the local
|
| 109 |
+
Transformers installation also provides native RWKV-7 support.
|
| 110 |
+
|
| 111 |
+
```python
|
| 112 |
+
import torch
|
| 113 |
+
from transformers import (
|
| 114 |
+
AutoModelForCausalLM,
|
| 115 |
+
AutoTokenizer,
|
| 116 |
+
PreTrainedConfig,
|
| 117 |
+
)
|
| 118 |
+
|
| 119 |
+
model_id = "aabbdev/RWKV7-1.5B-20260805"
|
| 120 |
+
tokenizer = AutoTokenizer.from_pretrained(
|
| 121 |
+
model_id,
|
| 122 |
+
config=PreTrainedConfig(),
|
| 123 |
+
)
|
| 124 |
+
model = AutoModelForCausalLM.from_pretrained(
|
| 125 |
+
model_id,
|
| 126 |
+
trust_remote_code=True,
|
| 127 |
+
dtype=torch.bfloat16,
|
| 128 |
+
)
|
| 129 |
+
```
|
| 130 |
+
|
| 131 |
+
The recurrent cache returned by the model can be passed back for incremental
|
| 132 |
+
decoding. Use an `attention_mask` for padded batches.
|
| 133 |
+
|
| 134 |
+
## Chat quickstart
|
| 135 |
+
|
| 136 |
+
```python
|
| 137 |
+
import re
|
| 138 |
+
|
| 139 |
+
import torch
|
| 140 |
+
from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedConfig
|
| 141 |
+
|
| 142 |
+
|
| 143 |
+
THINK_RE = re.compile(r"\A<think>?\s*(.*?)\s*</think>?", re.DOTALL)
|
| 144 |
+
|
| 145 |
+
|
| 146 |
+
def assistant_content(completion, thinking, *, close_incomplete=False):
|
| 147 |
+
prefix = "<think" if thinking else "<think></think"
|
| 148 |
+
reply = prefix + completion
|
| 149 |
+
thinking_block = THINK_RE.match(reply)
|
| 150 |
+
if thinking:
|
| 151 |
+
if thinking_block is not None or not close_incomplete:
|
| 152 |
+
return reply.strip()
|
| 153 |
+
return f"{reply.rstrip()}\n</think>".strip()
|
| 154 |
+
return "" if thinking_block is None else reply[thinking_block.end():].strip()
|
| 155 |
+
|
| 156 |
+
model_id = "aabbdev/RWKV7-1.5B-20260805"
|
| 157 |
+
tokenizer = AutoTokenizer.from_pretrained(
|
| 158 |
+
model_id,
|
| 159 |
+
config=PreTrainedConfig(),
|
| 160 |
+
)
|
| 161 |
+
model = AutoModelForCausalLM.from_pretrained(
|
| 162 |
+
model_id,
|
| 163 |
+
trust_remote_code=True,
|
| 164 |
+
dtype=torch.bfloat16,
|
| 165 |
+
).to("cuda")
|
| 166 |
+
|
| 167 |
+
messages = [{"role": "user", "content": "Explain why RWKV uses constant state."}]
|
| 168 |
+
thinking = False
|
| 169 |
+
max_new_tokens = 256
|
| 170 |
+
inputs = tokenizer.apply_chat_template(
|
| 171 |
+
messages,
|
| 172 |
+
tokenize=True,
|
| 173 |
+
add_generation_prompt=True,
|
| 174 |
+
thinking=thinking,
|
| 175 |
+
return_dict=True,
|
| 176 |
+
return_tensors="pt",
|
| 177 |
+
).to(model.device)
|
| 178 |
+
|
| 179 |
+
output = model.generate(
|
| 180 |
+
**inputs,
|
| 181 |
+
max_new_tokens=max_new_tokens,
|
| 182 |
+
do_sample=True,
|
| 183 |
+
temperature=1.0,
|
| 184 |
+
top_p=0.5,
|
| 185 |
+
eos_token_id=0,
|
| 186 |
+
pad_token_id=0,
|
| 187 |
+
stop_strings=["\n\nUser:"],
|
| 188 |
+
tokenizer=tokenizer,
|
| 189 |
+
)
|
| 190 |
+
completion = tokenizer.decode(
|
| 191 |
+
output[0, inputs["input_ids"].shape[1]:],
|
| 192 |
+
skip_special_tokens=True,
|
| 193 |
+
)
|
| 194 |
+
completion = completion.split("\n\nUser:", 1)[0]
|
| 195 |
+
reached_token_limit = output.shape[1] - inputs["input_ids"].shape[1] >= max_new_tokens
|
| 196 |
+
print(
|
| 197 |
+
assistant_content(
|
| 198 |
+
completion,
|
| 199 |
+
thinking,
|
| 200 |
+
close_incomplete=reached_token_limit,
|
| 201 |
+
)
|
| 202 |
+
)
|
| 203 |
+
```
|
| 204 |
+
|
| 205 |
+
Set `thinking=True` for the RWKV thinking prefix. The intentional generation
|
| 206 |
+
prefixes are `Assistant: <think></think` and `Assistant: <think`; do not append a
|
| 207 |
+
closing `>` to them. The post-processing above reconstructs that prefix before
|
| 208 |
+
removing an empty thinking block or preserving an enabled one. If generation hits
|
| 209 |
+
the token limit inside thinking, it closes the displayed block before returning it.
|
| 210 |
+
Reference stops are token ID `0` and `\n\nUser:`.
|
| 211 |
+
|
| 212 |
+
Strip trailing spaces from user input. The official RWKV prompt guide is available
|
| 213 |
+
in [`RWKV7-G1x-templates.txt`](https://github.com/BlinkDL/RWKV-LM/blob/main/RWKV-v7/RWKV7-G1x-templates.txt).
|
| 214 |
+
|
| 215 |
+
## Optimized local inference
|
| 216 |
+
|
| 217 |
+
Install the versions listed in `inference/requirements.txt`, then run the bundled
|
| 218 |
+
interactive chat:
|
| 219 |
+
|
| 220 |
+
```bash
|
| 221 |
+
python inference/generate.py --model aabbdev/RWKV7-1.5B-20260805 --backend auto --interactive
|
| 222 |
+
```
|
| 223 |
+
|
| 224 |
+
Or independent prompts separated by blank lines:
|
| 225 |
+
|
| 226 |
+
```bash
|
| 227 |
+
python inference/generate.py \
|
| 228 |
+
--model aabbdev/RWKV7-1.5B-20260805 \
|
| 229 |
+
--backend auto \
|
| 230 |
+
--input-file prompts.txt
|
| 231 |
+
```
|
| 232 |
+
|
| 233 |
+
`--backend auto` uses validated exact optimized boundaries and otherwise falls
|
| 234 |
+
back to PyTorch. Full explicit TileLang execution can change floating-point
|
| 235 |
+
operation order and requires checkpoint-, dtype-, shape-, and device-specific
|
| 236 |
+
parity validation.
|
| 237 |
+
|
| 238 |
+
## Tokenizer
|
| 239 |
+
|
| 240 |
+
The model root contains one self-contained tokenizer artifact: `tokenizer.json`.
|
| 241 |
+
Textual `vocab.json` and `rwkv_vocab_v20230424.txt` files are intentionally omitted
|
| 242 |
+
because they would duplicate the tokenizer used by Transformers. The tokenizer is
|
| 243 |
+
loaded natively as `PreTrainedTokenizerFast` and never executes remote Python code.
|
| 244 |
+
The explicit generic config prevents `AutoTokenizer` from probing the remote model
|
| 245 |
+
configuration and emitting a harmless model-type fallback warning.
|
| 246 |
+
|
| 247 |
+
## Intended use and limitations
|
| 248 |
+
|
| 249 |
+
- This is a base causal language model. Quality, instruction following, and
|
| 250 |
+
language behavior depend on the checkpoint and downstream prompting or
|
| 251 |
+
post-training.
|
| 252 |
+
- Assisted or speculative decoding that requires recurrent-cache rollback is not
|
| 253 |
+
supported without retaining prior state snapshots.
|
| 254 |
+
- Optimized support depends on GPU architecture, dtype, batch, and shape.
|
| 255 |
+
Unsupported `auto` configurations fall back to pure PyTorch.
|
| 256 |
+
- Explicit full TileLang execution can change floating-point operation order and
|
| 257 |
+
requires checkpoint-, dtype-, shape-, and device-specific parity validation.
|
| 258 |
+
- No safety, bias, toxicity, factuality, or high-stakes-use evaluation is claimed
|
| 259 |
+
by this model card.
|
| 260 |
+
|
| 261 |
+
## License and provenance
|
| 262 |
+
|
| 263 |
+
The model weights use the locked profile license `apache-2.0`. The exported inference bundle is licensed separately under [Apache-2.0](LICENSE). The bundled Transformers configuration and modeling modules
|
| 264 |
+
retain their Apache-2.0 headers. See [`NOTICE`](NOTICE) and the source checkpoint
|
| 265 |
+
link above for provenance.
|
| 266 |
+
|
| 267 |
+
## Citation
|
| 268 |
+
|
| 269 |
+
```bibtex
|
| 270 |
+
@misc{peng2025250314456,
|
| 271 |
+
title = {RWKV-7 "Goose" with Expressive Dynamic State Evolution},
|
| 272 |
+
author = {Bo Peng and Ruichong Zhang and Daniel Goldstein and Eric Alcaide and Xingjian Du and Haowen Hou and Jiaju Lin and Jiaxing Liu and Janna Lu and William Merrill and Guangyu Song and Kaifeng Tan and Saiteja Utpala and Nathan Wilce and Johan S. Wind and Tianyi Wu and Daniel Wuttke and Christian Zhou-Zheng},
|
| 273 |
+
year = {2025},
|
| 274 |
+
eprint = {2503.14456v2},
|
| 275 |
+
archivePrefix = {arXiv},
|
| 276 |
+
primaryClass = {cs.CL},
|
| 277 |
+
url = {https://arxiv.org/abs/2503.14456v2},
|
| 278 |
+
}
|
| 279 |
+
```
|
chat_template.jinja
ADDED
|
@@ -0,0 +1,56 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{%- set add_generation_prompt = add_generation_prompt | default(true) -%}
|
| 2 |
+
{%- set thinking = thinking | default(false) -%}
|
| 3 |
+
{%- set bos_token = bos_token | default('', true) -%}
|
| 4 |
+
{%- set tools = tools | default([], true) -%}
|
| 5 |
+
{%- set ns = namespace(system_prompt='') -%}
|
| 6 |
+
{%- for message in messages -%}
|
| 7 |
+
{%- if message.role == 'system' -%}
|
| 8 |
+
{%- set ns.system_prompt = message.content | trim -%}
|
| 9 |
+
{%- endif -%}
|
| 10 |
+
{%- endfor -%}
|
| 11 |
+
{{- bos_token -}}
|
| 12 |
+
{%- if ns.system_prompt or tools | length > 0 -%}
|
| 13 |
+
{{ 'System: ' }}{{ ns.system_prompt }}
|
| 14 |
+
{%- if tools | length > 0 -%}
|
| 15 |
+
{%- if ns.system_prompt %}{{ '\n' }}{%- endif -%}
|
| 16 |
+
{{ 'Tools:\n' -}}
|
| 17 |
+
{{ tools | tojson }}
|
| 18 |
+
{{ '\nWhen using a tool, return only a compact JSON function call in a ```json block, like {"name":"calculator","arguments":{"expression":"2+2"}}. The `name` field must be top-level, never inside `arguments`. Do not copy the tool schema into arguments. Otherwise answer normally.' }}
|
| 19 |
+
{%- endif -%}
|
| 20 |
+
{{ '\n\n' }}
|
| 21 |
+
{%- endif -%}
|
| 22 |
+
{%- for message in messages -%}
|
| 23 |
+
{%- if message.role == 'user' -%}
|
| 24 |
+
{{ 'User: ' ~ (message.content | trim) ~ '\n\n' }}
|
| 25 |
+
{%- elif message.role == 'assistant' -%}
|
| 26 |
+
{%- set content = message.content | default('', true) | trim -%}
|
| 27 |
+
{{ 'Assistant:' }}
|
| 28 |
+
{%- if message.tool_calls is defined and message.tool_calls | length > 0 -%}
|
| 29 |
+
{%- if content %}{{ ' ' ~ content ~ '\n' }}{%- endif -%}
|
| 30 |
+
{%- for tool_call in message.tool_calls -%}
|
| 31 |
+
{%- if tool_call.function is defined -%}
|
| 32 |
+
{%- set name = tool_call.function.name | default('') -%}
|
| 33 |
+
{%- set args = tool_call.function.arguments | default({}, true) -%}
|
| 34 |
+
{%- else -%}
|
| 35 |
+
{%- set name = tool_call.name | default('') -%}
|
| 36 |
+
{%- set args = tool_call.arguments | default({}, true) -%}
|
| 37 |
+
{%- endif -%}
|
| 38 |
+
{{ ' ```json\n' -}}
|
| 39 |
+
{{ '{"name": ' }}{{ name | tojson }}{{ ', "arguments": ' }}{% if args is string %}{{ args }}{% else %}{{ args | tojson }}{% endif %}{{ '}\n' -}}
|
| 40 |
+
{{ '```' }}{{ '\n' if not loop.last else '' }}
|
| 41 |
+
{%- endfor -%}
|
| 42 |
+
{%- elif content -%}
|
| 43 |
+
{{ ' ' ~ content }}
|
| 44 |
+
{%- endif -%}
|
| 45 |
+
{{ '\n\n' }}
|
| 46 |
+
{%- elif message.role == 'tool' -%}
|
| 47 |
+
{{ 'User: Function output:\n' ~ (message.content | trim) ~ '\n\n' }}
|
| 48 |
+
{%- endif -%}
|
| 49 |
+
{%- endfor -%}
|
| 50 |
+
{%- if add_generation_prompt -%}
|
| 51 |
+
{%- if thinking -%}
|
| 52 |
+
{{ 'Assistant: <think' }}
|
| 53 |
+
{%- else -%}
|
| 54 |
+
{{ 'Assistant: <think></think' }}
|
| 55 |
+
{%- endif -%}
|
| 56 |
+
{%- endif -%}
|
config.json
ADDED
|
@@ -0,0 +1,31 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"a_low_rank_dim": 96,
|
| 3 |
+
"architectures": [
|
| 4 |
+
"Rwkv7ForCausalLM"
|
| 5 |
+
],
|
| 6 |
+
"auto_map": {
|
| 7 |
+
"AutoConfig": "configuration_rwkv7.Rwkv7Config",
|
| 8 |
+
"AutoModel": "modeling_rwkv7.Rwkv7Model",
|
| 9 |
+
"AutoModelForCausalLM": "modeling_rwkv7.Rwkv7ForCausalLM"
|
| 10 |
+
},
|
| 11 |
+
"bos_token_id": 0,
|
| 12 |
+
"decay_low_rank_dim": 96,
|
| 13 |
+
"dtype": "bfloat16",
|
| 14 |
+
"eos_token_id": 0,
|
| 15 |
+
"gate_low_rank_dim": 256,
|
| 16 |
+
"head_dim": 64,
|
| 17 |
+
"hidden_size": 2048,
|
| 18 |
+
"intermediate_size": 8192,
|
| 19 |
+
"model_type": "rwkv7",
|
| 20 |
+
"norm_bias": true,
|
| 21 |
+
"norm_eps": 1e-05,
|
| 22 |
+
"num_heads": 32,
|
| 23 |
+
"num_hidden_layers": 24,
|
| 24 |
+
"pad_token_id": 0,
|
| 25 |
+
"tie_word_embeddings": false,
|
| 26 |
+
"use_cache": true,
|
| 27 |
+
"v_low_rank_dim": 64,
|
| 28 |
+
"vocab_size": 65536,
|
| 29 |
+
"wkv_implementation": "eager",
|
| 30 |
+
"wkv_state_dtype": "float32"
|
| 31 |
+
}
|
configuration_rwkv7.py
ADDED
|
@@ -0,0 +1,153 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright 2026 The RWKV team and The HuggingFace Inc. team. All rights reserved.
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
"""RWKV-7 (Goose) model configuration."""
|
| 15 |
+
|
| 16 |
+
from huggingface_hub.dataclasses import strict
|
| 17 |
+
|
| 18 |
+
from transformers.configuration_utils import PreTrainedConfig
|
| 19 |
+
from transformers.utils import auto_docstring
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
@auto_docstring(
|
| 23 |
+
checkpoint="RWKV/RWKV7-1.5B-20260805",
|
| 24 |
+
custom_intro="""
|
| 25 |
+
Configuration for [`Rwkv7Model`], an all-recurrent (attention-free) RWKV-7 "Goose"
|
| 26 |
+
model. Instantiating with the defaults yields the ~0.1B RWKV-7 configuration.
|
| 27 |
+
|
| 28 |
+
Parameter names follow the upstream RWKV reference implementation
|
| 29 |
+
(`BlinkDL/RWKV-LM`) rather than a renamed variant.
|
| 30 |
+
""",
|
| 31 |
+
)
|
| 32 |
+
@strict
|
| 33 |
+
class Rwkv7Config(PreTrainedConfig):
|
| 34 |
+
r"""
|
| 35 |
+
vocab_size (`int`, *optional*, defaults to 65536):
|
| 36 |
+
Vocabulary size (RWKV "world" tokenizer).
|
| 37 |
+
hidden_size (`int`, *optional*, defaults to 768):
|
| 38 |
+
Model width `C`.
|
| 39 |
+
num_hidden_layers (`int`, *optional*, defaults to 12):
|
| 40 |
+
Number of blocks.
|
| 41 |
+
head_dim (`int`, *optional*, defaults to 64):
|
| 42 |
+
Width of one WKV head. `hidden_size` must be divisible by it.
|
| 43 |
+
num_heads (`int`, *optional*, defaults to 12):
|
| 44 |
+
Number of WKV heads; must equal `hidden_size // head_dim`.
|
| 45 |
+
decay_low_rank_dim (`int`, *optional*):
|
| 46 |
+
Rank of the decay (`w`) LoRA. When omitted, follows the RWKV-LM recipe:
|
| 47 |
+
`max(32, round(2.5 * sqrt(hidden_size) / 32) * 32)`.
|
| 48 |
+
a_low_rank_dim (`int`, *optional*):
|
| 49 |
+
Rank of the in-context-learning-rate (`a`) LoRA. Uses the same default
|
| 50 |
+
formula as `decay_low_rank_dim`.
|
| 51 |
+
v_low_rank_dim (`int`, *optional*):
|
| 52 |
+
Rank of the value-residual (`v`) LoRA. Defaults to
|
| 53 |
+
`max(32, round(1.7 * sqrt(hidden_size) / 32) * 32)`. Unused on layer 0,
|
| 54 |
+
which *produces* `v_first` instead of mixing towards it.
|
| 55 |
+
gate_low_rank_dim (`int`, *optional*):
|
| 56 |
+
Rank of the output-gate (`g`) LoRA. Defaults to
|
| 57 |
+
`max(32, round(5 * sqrt(hidden_size) / 32) * 32)`.
|
| 58 |
+
intermediate_size (`int`, *optional*):
|
| 59 |
+
Channel-mix inner width. Defaults to `4 * hidden_size`.
|
| 60 |
+
norm_eps (`float`, *optional*, defaults to 1e-05):
|
| 61 |
+
Epsilon of every LayerNorm/GroupNorm in the model.
|
| 62 |
+
norm_bias (`bool`, *optional*, defaults to `True`):
|
| 63 |
+
Whether the norms carry a bias.
|
| 64 |
+
tie_word_embeddings (`bool`, *optional*, defaults to `False`):
|
| 65 |
+
Whether to tie the input embedding and the LM head.
|
| 66 |
+
use_cache (`bool`, *optional*, defaults to `True`):
|
| 67 |
+
Whether to return the recurrent state.
|
| 68 |
+
wkv_state_dtype (`str`, *optional*, defaults to `"float32"`):
|
| 69 |
+
Precision the recurrent WKV state is carried and accumulated in,
|
| 70 |
+
independently of the activation dtype. The recurrence is unrolled over
|
| 71 |
+
the whole sequence, so a narrow state drifts; `"float32"` with fp16
|
| 72 |
+
activations is the combination the reference implementation uses.
|
| 73 |
+
`"float16"`/`"bfloat16"` trade that for a smaller state.
|
| 74 |
+
wkv_implementation (`str`, *optional*, defaults to `"eager"`):
|
| 75 |
+
Which WKV recurrence to use, by name, from
|
| 76 |
+
`models.rwkv7.modeling_rwkv7.RWKV7_WKV_FUNCTIONS`. `"eager"` is the
|
| 77 |
+
exact portable PyTorch path and preserves the reference token order for
|
| 78 |
+
both prefill and decode. `"chunked"` is an explicit faster
|
| 79 |
+
chunk-parallel rearrangement whose floating-point order differs. Register
|
| 80 |
+
an entry in that mapping to plug in a fused or varlen kernel without
|
| 81 |
+
forking the model.
|
| 82 |
+
bos_token_id (`int`, *optional*, defaults to 0):
|
| 83 |
+
Beginning-of-sequence id. The RWKV world tokenizer has no dedicated BOS
|
| 84 |
+
token and the reference implementation prepends nothing, so this exists to
|
| 85 |
+
satisfy `GenerationMixin` rather than to be emitted.
|
| 86 |
+
eos_token_id (`int`, *optional*, defaults to 0):
|
| 87 |
+
End-of-sequence id, id 0 in the RWKV world vocabulary.
|
| 88 |
+
pad_token_id (`int`, *optional*, defaults to 0):
|
| 89 |
+
Padding id, the same id 0. Set deliberately rather than left `None`:
|
| 90 |
+
`generate` needs one to pad a batch, and without it a batched call either
|
| 91 |
+
raised or fell back to the eos id with a warning on every step.
|
| 92 |
+
|
| 93 |
+
```python
|
| 94 |
+
>>> from transformers import Rwkv7Config, Rwkv7Model
|
| 95 |
+
|
| 96 |
+
>>> configuration = Rwkv7Config()
|
| 97 |
+
>>> model = Rwkv7Model(configuration)
|
| 98 |
+
>>> configuration = model.config
|
| 99 |
+
```"""
|
| 100 |
+
|
| 101 |
+
model_type = "rwkv7"
|
| 102 |
+
keys_to_ignore_at_inference = ["state"]
|
| 103 |
+
|
| 104 |
+
vocab_size: int = 65536
|
| 105 |
+
hidden_size: int = 768
|
| 106 |
+
num_hidden_layers: int = 12
|
| 107 |
+
head_dim: int = 64
|
| 108 |
+
num_heads: int = 12
|
| 109 |
+
decay_low_rank_dim: int | None = None
|
| 110 |
+
a_low_rank_dim: int | None = None
|
| 111 |
+
v_low_rank_dim: int | None = None
|
| 112 |
+
gate_low_rank_dim: int | None = None
|
| 113 |
+
# `None` rather than a number: a literal default is correct for the default
|
| 114 |
+
# `hidden_size` and silently wrong for every other one, so a config built as
|
| 115 |
+
# `Rwkv7Config(hidden_size=4096, num_heads=64)` would come back with a channel-mix
|
| 116 |
+
# four times narrower than the architecture it names. `__post_init__` resolves it,
|
| 117 |
+
# and the resolved value is written to `config.json` either way.
|
| 118 |
+
intermediate_size: int | None = None
|
| 119 |
+
norm_eps: float = 1e-5
|
| 120 |
+
norm_bias: bool = True
|
| 121 |
+
tie_word_embeddings: bool = False
|
| 122 |
+
use_cache: bool = True
|
| 123 |
+
wkv_state_dtype: str = "float32"
|
| 124 |
+
wkv_implementation: str = "eager"
|
| 125 |
+
bos_token_id: int | None = 0
|
| 126 |
+
eos_token_id: int | None = 0
|
| 127 |
+
pad_token_id: int | None = 0
|
| 128 |
+
|
| 129 |
+
def __post_init__(self, **kwargs):
|
| 130 |
+
rank_unit = self.hidden_size**0.5 / 32
|
| 131 |
+
if self.decay_low_rank_dim is None:
|
| 132 |
+
self.decay_low_rank_dim = max(32, round(2.5 * rank_unit) * 32)
|
| 133 |
+
if self.a_low_rank_dim is None:
|
| 134 |
+
self.a_low_rank_dim = max(32, round(2.5 * rank_unit) * 32)
|
| 135 |
+
if self.v_low_rank_dim is None:
|
| 136 |
+
self.v_low_rank_dim = max(32, round(1.7 * rank_unit) * 32)
|
| 137 |
+
if self.gate_low_rank_dim is None:
|
| 138 |
+
self.gate_low_rank_dim = max(32, round(5 * rank_unit) * 32)
|
| 139 |
+
if self.intermediate_size is None:
|
| 140 |
+
self.intermediate_size = 4 * self.hidden_size
|
| 141 |
+
if self.wkv_state_dtype not in ("float32", "float16", "bfloat16"):
|
| 142 |
+
raise ValueError(f"wkv_state_dtype must be float32/float16/bfloat16, got {self.wkv_state_dtype}")
|
| 143 |
+
if self.hidden_size % self.head_dim != 0:
|
| 144 |
+
raise ValueError(f"hidden_size {self.hidden_size} must be divisible by head_dim {self.head_dim}")
|
| 145 |
+
if self.num_heads != self.hidden_size // self.head_dim:
|
| 146 |
+
raise ValueError(
|
| 147 |
+
f"num_heads must be hidden_size // head_dim = {self.hidden_size // self.head_dim}, "
|
| 148 |
+
f"got {self.num_heads}"
|
| 149 |
+
)
|
| 150 |
+
super().__post_init__(**kwargs)
|
| 151 |
+
|
| 152 |
+
|
| 153 |
+
__all__ = ["Rwkv7Config"]
|
generation_config.json
ADDED
|
@@ -0,0 +1,6 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"bos_token_id": 0,
|
| 3 |
+
"eos_token_id": 0,
|
| 4 |
+
"pad_token_id": 0,
|
| 5 |
+
"use_cache": true
|
| 6 |
+
}
|
inference/generate.py
ADDED
|
@@ -0,0 +1,239 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
import argparse
|
| 4 |
+
import importlib
|
| 5 |
+
import os
|
| 6 |
+
import re
|
| 7 |
+
import sys
|
| 8 |
+
from types import ModuleType
|
| 9 |
+
from pathlib import Path
|
| 10 |
+
from typing import Any
|
| 11 |
+
|
| 12 |
+
import torch
|
| 13 |
+
from transformers import StoppingCriteria, StoppingCriteriaList
|
| 14 |
+
|
| 15 |
+
if __package__ in {None, ""}:
|
| 16 |
+
package_name = "_rwkv7_release_inference"
|
| 17 |
+
package = ModuleType(package_name)
|
| 18 |
+
package.__package__ = package_name
|
| 19 |
+
package.__path__ = [str(Path(__file__).resolve().parent)]
|
| 20 |
+
sys.modules[package_name] = package
|
| 21 |
+
load_model_and_tokenizer = importlib.import_module(
|
| 22 |
+
f"{package_name}.model_loader"
|
| 23 |
+
).load_model_and_tokenizer
|
| 24 |
+
else:
|
| 25 |
+
from .model_loader import load_model_and_tokenizer
|
| 26 |
+
|
| 27 |
+
|
| 28 |
+
DTYPES = {
|
| 29 |
+
"bfloat16": torch.bfloat16,
|
| 30 |
+
"float16": torch.float16,
|
| 31 |
+
"float32": torch.float32,
|
| 32 |
+
}
|
| 33 |
+
STOP_TEXT = "\n\nUser:"
|
| 34 |
+
THINK_RE = re.compile(r"\A<think>?\s*(.*?)\s*</think>?", re.DOTALL)
|
| 35 |
+
|
| 36 |
+
|
| 37 |
+
class StopOnText(StoppingCriteria):
|
| 38 |
+
def __init__(self, tokenizer: Any, prompt_length: int, stop_text: str) -> None:
|
| 39 |
+
self.tokenizer = tokenizer
|
| 40 |
+
self.prompt_length = prompt_length
|
| 41 |
+
self.stop_text = stop_text
|
| 42 |
+
self.matched = False
|
| 43 |
+
|
| 44 |
+
def __call__(
|
| 45 |
+
self,
|
| 46 |
+
input_ids: torch.LongTensor,
|
| 47 |
+
scores: torch.FloatTensor,
|
| 48 |
+
**kwargs: Any,
|
| 49 |
+
) -> bool:
|
| 50 |
+
del scores, kwargs
|
| 51 |
+
completion = self.tokenizer.decode(
|
| 52 |
+
input_ids[0, self.prompt_length :],
|
| 53 |
+
skip_special_tokens=False,
|
| 54 |
+
)
|
| 55 |
+
self.matched = self.stop_text in completion
|
| 56 |
+
return self.matched
|
| 57 |
+
|
| 58 |
+
|
| 59 |
+
def _assistant_content(
|
| 60 |
+
completion: str, thinking: bool, *, close_incomplete: bool = False
|
| 61 |
+
) -> str:
|
| 62 |
+
prefix = "<think" if thinking else "<think></think"
|
| 63 |
+
reply = prefix + completion
|
| 64 |
+
thinking_block = THINK_RE.match(reply)
|
| 65 |
+
if thinking:
|
| 66 |
+
if thinking_block is not None or not close_incomplete:
|
| 67 |
+
return reply.strip()
|
| 68 |
+
return f"{reply.rstrip()}\n</think>".strip()
|
| 69 |
+
return "" if thinking_block is None else reply[thinking_block.end() :].strip()
|
| 70 |
+
|
| 71 |
+
|
| 72 |
+
def _prompt_ids(tokenizer: Any, messages: list[dict[str, str]], thinking: bool):
|
| 73 |
+
tokens = tokenizer.apply_chat_template(
|
| 74 |
+
messages,
|
| 75 |
+
tokenize=True,
|
| 76 |
+
add_generation_prompt=True,
|
| 77 |
+
thinking=thinking,
|
| 78 |
+
return_tensors="pt",
|
| 79 |
+
)
|
| 80 |
+
if hasattr(tokens, "input_ids"):
|
| 81 |
+
tokens = tokens.input_ids
|
| 82 |
+
elif isinstance(tokens, dict):
|
| 83 |
+
tokens = tokens["input_ids"]
|
| 84 |
+
if tokens.ndim == 1:
|
| 85 |
+
tokens = tokens.unsqueeze(0)
|
| 86 |
+
return tokens
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
@torch.inference_mode()
|
| 90 |
+
def generate_completion(
|
| 91 |
+
model: Any,
|
| 92 |
+
tokenizer: Any,
|
| 93 |
+
messages: list[dict[str, str]],
|
| 94 |
+
*,
|
| 95 |
+
device: str,
|
| 96 |
+
max_new_tokens: int,
|
| 97 |
+
temperature: float,
|
| 98 |
+
top_p: float,
|
| 99 |
+
thinking: bool,
|
| 100 |
+
) -> str:
|
| 101 |
+
input_ids = _prompt_ids(tokenizer, messages, thinking).to(device)
|
| 102 |
+
prompt_length = input_ids.shape[1]
|
| 103 |
+
stop_on_text = StopOnText(tokenizer, prompt_length, STOP_TEXT)
|
| 104 |
+
generation: dict[str, Any] = {
|
| 105 |
+
"input_ids": input_ids,
|
| 106 |
+
"attention_mask": torch.ones_like(input_ids),
|
| 107 |
+
"max_new_tokens": max_new_tokens,
|
| 108 |
+
"do_sample": temperature > 0,
|
| 109 |
+
"eos_token_id": 0,
|
| 110 |
+
"pad_token_id": 0,
|
| 111 |
+
"stopping_criteria": StoppingCriteriaList([stop_on_text]),
|
| 112 |
+
}
|
| 113 |
+
if temperature > 0:
|
| 114 |
+
generation["temperature"] = temperature
|
| 115 |
+
generation["top_p"] = top_p
|
| 116 |
+
output = model.generate(**generation)
|
| 117 |
+
completion_ids = output[0, prompt_length:]
|
| 118 |
+
completion = tokenizer.decode(completion_ids, skip_special_tokens=True)
|
| 119 |
+
if STOP_TEXT in completion:
|
| 120 |
+
completion = completion.split(STOP_TEXT, 1)[0]
|
| 121 |
+
reached_token_limit = (
|
| 122 |
+
completion_ids.shape[0] >= max_new_tokens and not stop_on_text.matched
|
| 123 |
+
)
|
| 124 |
+
return _assistant_content(
|
| 125 |
+
completion,
|
| 126 |
+
thinking,
|
| 127 |
+
close_incomplete=reached_token_limit,
|
| 128 |
+
)
|
| 129 |
+
|
| 130 |
+
|
| 131 |
+
def _interactive(
|
| 132 |
+
model: Any,
|
| 133 |
+
tokenizer: Any,
|
| 134 |
+
args: argparse.Namespace,
|
| 135 |
+
) -> None:
|
| 136 |
+
messages: list[dict[str, str]] = []
|
| 137 |
+
print("RWKV-7 Goose — /clear resets the conversation, /exit quits.")
|
| 138 |
+
while True:
|
| 139 |
+
try:
|
| 140 |
+
prompt = input(">>> ")
|
| 141 |
+
except EOFError:
|
| 142 |
+
break
|
| 143 |
+
if prompt == "/exit":
|
| 144 |
+
break
|
| 145 |
+
if prompt == "/clear":
|
| 146 |
+
messages.clear()
|
| 147 |
+
continue
|
| 148 |
+
prompt = prompt.strip()
|
| 149 |
+
if not prompt:
|
| 150 |
+
continue
|
| 151 |
+
messages.append({"role": "user", "content": prompt})
|
| 152 |
+
completion = generate_completion(
|
| 153 |
+
model,
|
| 154 |
+
tokenizer,
|
| 155 |
+
messages,
|
| 156 |
+
device=args.device,
|
| 157 |
+
max_new_tokens=args.max_new_tokens,
|
| 158 |
+
temperature=args.temperature,
|
| 159 |
+
top_p=args.top_p,
|
| 160 |
+
thinking=args.thinking,
|
| 161 |
+
)
|
| 162 |
+
print(completion)
|
| 163 |
+
messages.append({"role": "assistant", "content": completion})
|
| 164 |
+
|
| 165 |
+
|
| 166 |
+
def _file_prompts(
|
| 167 |
+
model: Any,
|
| 168 |
+
tokenizer: Any,
|
| 169 |
+
args: argparse.Namespace,
|
| 170 |
+
) -> None:
|
| 171 |
+
text = Path(args.input_file).read_text(encoding="utf-8")
|
| 172 |
+
prompts = [prompt.strip() for prompt in re.split(r"\n\s*\n", text) if prompt.strip()]
|
| 173 |
+
if not prompts:
|
| 174 |
+
raise ValueError("input file contains no prompts")
|
| 175 |
+
for prompt in prompts:
|
| 176 |
+
completion = generate_completion(
|
| 177 |
+
model,
|
| 178 |
+
tokenizer,
|
| 179 |
+
[{"role": "user", "content": prompt}],
|
| 180 |
+
device=args.device,
|
| 181 |
+
max_new_tokens=args.max_new_tokens,
|
| 182 |
+
temperature=args.temperature,
|
| 183 |
+
top_p=args.top_p,
|
| 184 |
+
thinking=args.thinking,
|
| 185 |
+
)
|
| 186 |
+
print(f"Prompt: {prompt}")
|
| 187 |
+
print(f"Completion: {completion}")
|
| 188 |
+
print()
|
| 189 |
+
|
| 190 |
+
|
| 191 |
+
def parse_args() -> argparse.Namespace:
|
| 192 |
+
parser = argparse.ArgumentParser(description="Generate with RWKV-7 Goose")
|
| 193 |
+
parser.add_argument("--model", required=True, help="Hub repo ID or local model directory")
|
| 194 |
+
mode = parser.add_mutually_exclusive_group(required=True)
|
| 195 |
+
mode.add_argument("--interactive", action="store_true")
|
| 196 |
+
mode.add_argument("--input-file")
|
| 197 |
+
parser.add_argument("--device", default="cuda")
|
| 198 |
+
parser.add_argument("--dtype", choices=("auto", *DTYPES), default="auto")
|
| 199 |
+
parser.add_argument("--state-dtype", choices=DTYPES, default="float32")
|
| 200 |
+
parser.add_argument("--backend", choices=("auto", "torch", "tilelang"), default="auto")
|
| 201 |
+
parser.add_argument("--max-new-tokens", type=int, default=300)
|
| 202 |
+
parser.add_argument("--temperature", type=float, default=1.0)
|
| 203 |
+
parser.add_argument("--top-p", type=float, default=0.5)
|
| 204 |
+
parser.add_argument("--seed", type=int, default=33377335)
|
| 205 |
+
parser.add_argument("--thinking", action="store_true")
|
| 206 |
+
return parser.parse_args()
|
| 207 |
+
|
| 208 |
+
|
| 209 |
+
def main() -> None:
|
| 210 |
+
args = parse_args()
|
| 211 |
+
if int(os.getenv("WORLD_SIZE", "1")) != 1:
|
| 212 |
+
raise RuntimeError("the bundled runtime supports one process and one GPU")
|
| 213 |
+
if int(os.getenv("RANK", "0")) != 0 or int(os.getenv("LOCAL_RANK", "0")) != 0:
|
| 214 |
+
raise RuntimeError("RANK and LOCAL_RANK must be zero")
|
| 215 |
+
if args.max_new_tokens <= 0:
|
| 216 |
+
raise ValueError("max-new-tokens must be positive")
|
| 217 |
+
if args.temperature < 0:
|
| 218 |
+
raise ValueError("temperature must be non-negative")
|
| 219 |
+
if not 0 < args.top_p <= 1:
|
| 220 |
+
raise ValueError("top-p must be in (0, 1]")
|
| 221 |
+
torch.manual_seed(args.seed)
|
| 222 |
+
model, tokenizer = load_model_and_tokenizer(
|
| 223 |
+
args.model,
|
| 224 |
+
device=args.device,
|
| 225 |
+
dtype=None if args.dtype == "auto" else DTYPES[args.dtype],
|
| 226 |
+
backend=args.backend,
|
| 227 |
+
state_dtype=args.state_dtype,
|
| 228 |
+
)
|
| 229 |
+
model.set_kernel_backend(args.backend)
|
| 230 |
+
if args.backend == "tilelang":
|
| 231 |
+
model.prepare_inference_weights()
|
| 232 |
+
if args.interactive:
|
| 233 |
+
_interactive(model, tokenizer, args)
|
| 234 |
+
else:
|
| 235 |
+
_file_prompts(model, tokenizer, args)
|
| 236 |
+
|
| 237 |
+
|
| 238 |
+
if __name__ == "__main__":
|
| 239 |
+
main()
|
inference/kernel.py
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
inference/model_loader.py
ADDED
|
@@ -0,0 +1,178 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
import json
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
|
| 6 |
+
import torch
|
| 7 |
+
from huggingface_hub import snapshot_download
|
| 8 |
+
from safetensors.torch import load_file
|
| 9 |
+
from transformers import AutoTokenizer, PreTrainedConfig
|
| 10 |
+
|
| 11 |
+
from .runtime import RWKV7Config, RWKV7ForCausalLM
|
| 12 |
+
|
| 13 |
+
|
| 14 |
+
SQUEEZE_MARKERS = (
|
| 15 |
+
".x_",
|
| 16 |
+
".k_",
|
| 17 |
+
"att.r",
|
| 18 |
+
"att.w",
|
| 19 |
+
"att.v0",
|
| 20 |
+
"att.v1",
|
| 21 |
+
"att.v2",
|
| 22 |
+
"att.a",
|
| 23 |
+
"att.g",
|
| 24 |
+
)
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def _resolve_model(model: str) -> tuple[Path, bool]:
|
| 28 |
+
local = Path(model).expanduser()
|
| 29 |
+
if local.is_dir():
|
| 30 |
+
return local.resolve(), False
|
| 31 |
+
return Path(snapshot_download(model)), True
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
def _optimized_config(native: dict, backend: str, state_dtype: str) -> RWKV7Config:
|
| 35 |
+
return RWKV7Config(
|
| 36 |
+
vocab_size=native["vocab_size"],
|
| 37 |
+
hidden_size=native["hidden_size"],
|
| 38 |
+
num_hidden_layers=native["num_hidden_layers"],
|
| 39 |
+
head_size=native["head_dim"],
|
| 40 |
+
intermediate_size=native["intermediate_size"],
|
| 41 |
+
decay_lora_rank=native["decay_low_rank_dim"],
|
| 42 |
+
a_lora_rank=native["a_low_rank_dim"],
|
| 43 |
+
gate_lora_rank=native["gate_low_rank_dim"],
|
| 44 |
+
value_lora_rank=native["v_low_rank_dim"],
|
| 45 |
+
layer_norm_epsilon=native.get("norm_eps", 1e-5),
|
| 46 |
+
use_cache=native.get("use_cache", True),
|
| 47 |
+
kernel_backend=backend,
|
| 48 |
+
recurrent_state_dtype=state_dtype,
|
| 49 |
+
tie_word_embeddings=native.get("tie_word_embeddings", False),
|
| 50 |
+
bos_token_id=native.get("bos_token_id"),
|
| 51 |
+
eos_token_id=native.get("eos_token_id", 0),
|
| 52 |
+
pad_token_id=native.get("pad_token_id", 0),
|
| 53 |
+
)
|
| 54 |
+
|
| 55 |
+
|
| 56 |
+
def _native_key_to_optimized(key: str) -> str:
|
| 57 |
+
return key.removeprefix("rwkv7.")
|
| 58 |
+
|
| 59 |
+
|
| 60 |
+
def _native_tensor_to_optimized(key: str, tensor: torch.Tensor) -> torch.Tensor:
|
| 61 |
+
return tensor.squeeze() if any(marker in key for marker in SQUEEZE_MARKERS) else tensor
|
| 62 |
+
|
| 63 |
+
|
| 64 |
+
def _checked_weight(path: Path, *, hub_snapshot: bool) -> Path:
|
| 65 |
+
if path.is_symlink():
|
| 66 |
+
if not hub_snapshot:
|
| 67 |
+
raise RuntimeError(f"local safetensor must not be a symlink: {path.name}")
|
| 68 |
+
path = path.resolve(strict=True)
|
| 69 |
+
if not path.is_file():
|
| 70 |
+
raise RuntimeError(f"safetensor must be a regular file: {path.name}")
|
| 71 |
+
return path
|
| 72 |
+
|
| 73 |
+
|
| 74 |
+
def _weight_plan(
|
| 75 |
+
model_dir: Path, *, hub_snapshot: bool
|
| 76 |
+
) -> list[tuple[Path, set[str] | None]]:
|
| 77 |
+
present = sorted(model_dir.glob("model*.safetensors"))
|
| 78 |
+
if not present:
|
| 79 |
+
raise FileNotFoundError(f"no safetensors found in {model_dir}")
|
| 80 |
+
checked = {
|
| 81 |
+
path.name: _checked_weight(path, hub_snapshot=hub_snapshot) for path in present
|
| 82 |
+
}
|
| 83 |
+
index_path = model_dir / "model.safetensors.index.json"
|
| 84 |
+
if not index_path.is_file():
|
| 85 |
+
if [path.name for path in present] != ["model.safetensors"]:
|
| 86 |
+
raise RuntimeError("multiple safetensors require model.safetensors.index.json")
|
| 87 |
+
return [(checked["model.safetensors"], None)]
|
| 88 |
+
|
| 89 |
+
index = json.loads(index_path.read_text(encoding="utf-8"))
|
| 90 |
+
weight_map = index.get("weight_map")
|
| 91 |
+
if not isinstance(weight_map, dict) or not weight_map:
|
| 92 |
+
raise RuntimeError("model.safetensors.index.json has no weight_map")
|
| 93 |
+
keys_by_file: dict[str, set[str]] = {}
|
| 94 |
+
for key, filename in weight_map.items():
|
| 95 |
+
if not isinstance(key, str) or not isinstance(filename, str):
|
| 96 |
+
raise RuntimeError("invalid weight_map entry")
|
| 97 |
+
keys_by_file.setdefault(filename, set()).add(key)
|
| 98 |
+
if set(keys_by_file) != {path.name for path in present}:
|
| 99 |
+
raise RuntimeError("indexed and present safetensors files differ")
|
| 100 |
+
return [(checked[filename], keys_by_file[filename]) for filename in sorted(keys_by_file)]
|
| 101 |
+
|
| 102 |
+
|
| 103 |
+
def _load_native_model(
|
| 104 |
+
model_dir: Path,
|
| 105 |
+
native_config: dict,
|
| 106 |
+
backend: str,
|
| 107 |
+
state_dtype: str,
|
| 108 |
+
*,
|
| 109 |
+
hub_snapshot: bool,
|
| 110 |
+
) -> RWKV7ForCausalLM:
|
| 111 |
+
config = _optimized_config(native_config, backend, state_dtype)
|
| 112 |
+
with torch.device("meta"):
|
| 113 |
+
model = RWKV7ForCausalLM(config)
|
| 114 |
+
expected = set(model.state_dict())
|
| 115 |
+
seen: set[str] = set()
|
| 116 |
+
for weight_file, indexed_keys in _weight_plan(
|
| 117 |
+
model_dir, hub_snapshot=hub_snapshot
|
| 118 |
+
):
|
| 119 |
+
native_shard = load_file(weight_file, device="cpu")
|
| 120 |
+
if indexed_keys is not None and set(native_shard) != indexed_keys:
|
| 121 |
+
raise RuntimeError(f"tensor keys in {weight_file.name} do not match the index")
|
| 122 |
+
shard = {}
|
| 123 |
+
for native_key, tensor in native_shard.items():
|
| 124 |
+
optimized_key = _native_key_to_optimized(native_key)
|
| 125 |
+
if optimized_key in seen or optimized_key in shard:
|
| 126 |
+
raise RuntimeError(f"duplicate optimized tensor key: {optimized_key}")
|
| 127 |
+
shard[optimized_key] = _native_tensor_to_optimized(optimized_key, tensor)
|
| 128 |
+
unexpected = sorted(set(shard) - expected)
|
| 129 |
+
if unexpected:
|
| 130 |
+
raise RuntimeError(f"native checkpoint has unexpected keys: {unexpected}")
|
| 131 |
+
model.load_state_dict(shard, strict=False, assign=True)
|
| 132 |
+
seen.update(shard)
|
| 133 |
+
missing = sorted(expected - seen)
|
| 134 |
+
if missing:
|
| 135 |
+
raise RuntimeError(f"native checkpoint is missing optimized runtime keys: {missing}")
|
| 136 |
+
return model
|
| 137 |
+
|
| 138 |
+
|
| 139 |
+
def _load_tokenizer(model_dir: Path):
|
| 140 |
+
return AutoTokenizer.from_pretrained(
|
| 141 |
+
model_dir,
|
| 142 |
+
config=PreTrainedConfig(),
|
| 143 |
+
local_files_only=True,
|
| 144 |
+
)
|
| 145 |
+
|
| 146 |
+
|
| 147 |
+
def load_model_and_tokenizer(
|
| 148 |
+
model: str,
|
| 149 |
+
*,
|
| 150 |
+
device: str,
|
| 151 |
+
dtype: torch.dtype | None,
|
| 152 |
+
backend: str,
|
| 153 |
+
state_dtype: str,
|
| 154 |
+
):
|
| 155 |
+
model_dir, hub_snapshot = _resolve_model(model)
|
| 156 |
+
native_config = json.loads((model_dir / "config.json").read_text(encoding="utf-8"))
|
| 157 |
+
architectures = set(native_config.get("architectures", []))
|
| 158 |
+
if architectures != {"Rwkv7ForCausalLM"}:
|
| 159 |
+
raise ValueError(f"unsupported RWKV-7 architecture: {sorted(architectures)}")
|
| 160 |
+
loaded = _load_native_model(
|
| 161 |
+
model_dir,
|
| 162 |
+
native_config,
|
| 163 |
+
backend,
|
| 164 |
+
state_dtype,
|
| 165 |
+
hub_snapshot=hub_snapshot,
|
| 166 |
+
)
|
| 167 |
+
if dtype is None:
|
| 168 |
+
dtype_name = str(native_config.get("dtype", "bfloat16")).removeprefix("torch.")
|
| 169 |
+
try:
|
| 170 |
+
dtype = {
|
| 171 |
+
"bfloat16": torch.bfloat16,
|
| 172 |
+
"float16": torch.float16,
|
| 173 |
+
"float32": torch.float32,
|
| 174 |
+
}[dtype_name]
|
| 175 |
+
except KeyError as error:
|
| 176 |
+
raise ValueError(f"unsupported model dtype: {dtype_name}") from error
|
| 177 |
+
loaded = loaded.to(device=device, dtype=dtype).eval()
|
| 178 |
+
return loaded, _load_tokenizer(model_dir)
|
inference/requirements.txt
ADDED
|
@@ -0,0 +1,5 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
transformers>=5.15,<6
|
| 2 |
+
huggingface-hub>=0.34
|
| 3 |
+
safetensors>=0.5
|
| 4 |
+
jinja2>=3.1,<4
|
| 5 |
+
tilelang==0.1.12
|
inference/runtime.py
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
model.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:84ccbb857c84e00cefc48b233937ada79c411e491df25fb21aed23237f39a14f
|
| 3 |
+
size 3055418240
|
modeling_rwkv7.py
ADDED
|
@@ -0,0 +1,1186 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright 2026 The RWKV team and The HuggingFace Inc. team. All rights reserved.
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
"""PyTorch RWKV-7 ("Goose") model."""
|
| 15 |
+
|
| 16 |
+
import math
|
| 17 |
+
from dataclasses import dataclass
|
| 18 |
+
|
| 19 |
+
import torch
|
| 20 |
+
from torch import nn
|
| 21 |
+
|
| 22 |
+
from transformers import initialization as init
|
| 23 |
+
from transformers.cache_utils import Cache, LinearAttentionLayer
|
| 24 |
+
from transformers.generation import GenerationMixin
|
| 25 |
+
from transformers.modeling_layers import GradientCheckpointingLayer
|
| 26 |
+
from transformers.modeling_utils import PreTrainedModel
|
| 27 |
+
from transformers.utils import ModelOutput, auto_docstring, can_return_tuple, logging
|
| 28 |
+
from .configuration_rwkv7 import Rwkv7Config
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
logger = logging.get_logger(__name__)
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
def _rwkv7_lora_orthogonal_(tensor: torch.Tensor, scale: float = 0.1) -> None:
|
| 35 |
+
"""RWKV-LM compensates the gain when a low-rank output factor is tall."""
|
| 36 |
+
rows, columns = tensor.shape
|
| 37 |
+
gain = math.sqrt(rows / columns) if rows > columns else 1.0
|
| 38 |
+
init.orthogonal_(tensor, gain=gain * scale)
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
class Rwkv7TokenShift(nn.Module):
|
| 42 |
+
"""`prev_token(x)`: the previous token's hidden state, zero at sequence start.
|
| 43 |
+
|
| 44 |
+
Carries one vector per layer per stream across forward calls, which is what
|
| 45 |
+
makes incremental decoding exact rather than approximate.
|
| 46 |
+
"""
|
| 47 |
+
|
| 48 |
+
def forward(
|
| 49 |
+
self,
|
| 50 |
+
x: torch.Tensor,
|
| 51 |
+
shift_state: torch.Tensor | None,
|
| 52 |
+
cu_seq_lens: torch.Tensor | None = None,
|
| 53 |
+
keep: torch.Tensor | None = None,
|
| 54 |
+
):
|
| 55 |
+
# x: [batch, seq_len, hidden]; shift_state: [batch, hidden] or None
|
| 56 |
+
if shift_state is None:
|
| 57 |
+
prev = torch.zeros_like(x[:, :1])
|
| 58 |
+
else:
|
| 59 |
+
prev = shift_state[:, None]
|
| 60 |
+
shifted = torch.cat([prev, x[:, :-1]], dim=1)
|
| 61 |
+
if cu_seq_lens is not None:
|
| 62 |
+
# In a packed row the token before a segment's first one belongs to the
|
| 63 |
+
# PREVIOUS sequence. Resetting the recurrent state per segment does not
|
| 64 |
+
# cover this (the shift reaches back through it), so the first token of
|
| 65 |
+
# each segment gets the zero shift a sequence start is supposed to see.
|
| 66 |
+
positions = torch.arange(x.shape[1], device=x.device)
|
| 67 |
+
starts = (positions[:, None] == cu_seq_lens[None, :-1]).any(dim=1)
|
| 68 |
+
shifted = torch.where(starts[None, :, None], torch.zeros_like(shifted), shifted)
|
| 69 |
+
if keep is None:
|
| 70 |
+
return shifted, x[:, -1]
|
| 71 |
+
# The state handed back must be the last REAL token, not the last position.
|
| 72 |
+
# They coincide under left padding, which is why this went unnoticed; under
|
| 73 |
+
# right padding `x[:, -1]` is a blanked pad, so a continuation resumes from
|
| 74 |
+
# zero instead of from where the sequence actually got to.
|
| 75 |
+
# Integer arithmetic, not the activation dtype. `keep` arrives cast to the
|
| 76 |
+
# model's dtype, and bf16 carries eight mantissa bits: past position 256 the
|
| 77 |
+
# products are no longer distinct, `argmax` returns the first of a tie, and
|
| 78 |
+
# the state comes back from a token several places short of the last real one
|
| 79 |
+
# -- 1022 for 1023 at length 1024, 4088 for 4095 at 4096. An all-ones mask is
|
| 80 |
+
# enough to trigger it, which is what `generate` sends, so this was not
|
| 81 |
+
# confined to padded batches.
|
| 82 |
+
positions = torch.arange(x.shape[1], device=x.device)
|
| 83 |
+
last_real = ((keep.squeeze(-1) > 0) * positions).argmax(dim=-1)
|
| 84 |
+
return shifted, x[torch.arange(x.shape[0], device=x.device), last_real]
|
| 85 |
+
|
| 86 |
+
|
| 87 |
+
# e^-0.5. The decay LoRA emits w_log = -INV_SQRT_E * sigmoid(...), so w_log lies in
|
| 88 |
+
# (-e^-0.5, 0) and the per-step decay exp(w_log) lies in (exp(-e^-0.5), 1), i.e.
|
| 89 |
+
# (0.5452, 1). See RWKV-7 reference. Note the floor is exp(-e^-0.5), not e^-0.5: this
|
| 90 |
+
# comment said the latter, and `rwkv7_chunked` built its chunk_size bound on it.
|
| 91 |
+
_INV_SQRT_E = 0.6065306597126334
|
| 92 |
+
|
| 93 |
+
|
| 94 |
+
def _rwkv7_matmul(x: torch.Tensor, weight: torch.Tensor, tokenwise: bool) -> torch.Tensor:
|
| 95 |
+
"""Preserve decode GEMM shapes across a sequence when exact execution is requested."""
|
| 96 |
+
if not tokenwise or x.ndim < 3 or x.shape[1] == 1:
|
| 97 |
+
return x @ weight
|
| 98 |
+
return torch.stack([x[:, token_index] @ weight for token_index in range(x.shape[1])], dim=1)
|
| 99 |
+
|
| 100 |
+
|
| 101 |
+
def _rwkv7_linear(module: nn.Linear, x: torch.Tensor, tokenwise: bool) -> torch.Tensor:
|
| 102 |
+
if not tokenwise or x.ndim < 3 or x.shape[1] == 1:
|
| 103 |
+
return module(x)
|
| 104 |
+
return torch.stack([module(x[:, token_index]) for token_index in range(x.shape[1])], dim=1)
|
| 105 |
+
|
| 106 |
+
|
| 107 |
+
def rwkv7_recurrent(
|
| 108 |
+
r: torch.Tensor,
|
| 109 |
+
w_log: torch.Tensor,
|
| 110 |
+
k: torch.Tensor,
|
| 111 |
+
v: torch.Tensor,
|
| 112 |
+
kk: torch.Tensor,
|
| 113 |
+
a: torch.Tensor,
|
| 114 |
+
state: torch.Tensor,
|
| 115 |
+
compute_dtype: torch.dtype = torch.float32,
|
| 116 |
+
) -> tuple[torch.Tensor, torch.Tensor]:
|
| 117 |
+
"""Reference RWKV-7 WKV recurrence (generalised delta rule).
|
| 118 |
+
|
| 119 |
+
All inputs are `[batch, seq_len, num_heads, head_dim]`; `state` is
|
| 120 |
+
`[batch, num_heads, head_dim, head_dim]` and is carried across calls.
|
| 121 |
+
|
| 122 |
+
Per token, with `S` the per-head state::
|
| 123 |
+
|
| 124 |
+
decay = exp(w_log)
|
| 125 |
+
vk = v @ k
|
| 126 |
+
ab = (-kk) @ (kk * a)
|
| 127 |
+
S = S * decay + S @ ab + vk
|
| 128 |
+
out = S @ r
|
| 129 |
+
|
| 130 |
+
The state axes are (value, key), as in the reference implementation. Products
|
| 131 |
+
made from activations keep the activation dtype; the carried-state arithmetic
|
| 132 |
+
uses `compute_dtype`, and a narrower configured state is stored at every token
|
| 133 |
+
boundary before the output projection. Those boundaries are observable in a
|
| 134 |
+
recurrent BF16 model, so an algebraically transposed FP32 implementation is not
|
| 135 |
+
equivalent.
|
| 136 |
+
"""
|
| 137 |
+
batch, seq_len, num_heads, head_dim = r.shape
|
| 138 |
+
dtype = r.dtype
|
| 139 |
+
state_dtype = state.dtype
|
| 140 |
+
state = state.to(compute_dtype)
|
| 141 |
+
out = torch.empty(batch, seq_len, num_heads, head_dim, device=r.device, dtype=compute_dtype)
|
| 142 |
+
|
| 143 |
+
for t in range(seq_len):
|
| 144 |
+
decay = torch.exp(w_log[:, t])[:, :, None, :] # [B, H, 1, K]
|
| 145 |
+
kk_t = kk[:, t]
|
| 146 |
+
value_key = v[:, t][..., None] @ k[:, t][:, :, None, :]
|
| 147 |
+
anti_value_key = (-kk_t)[..., None] @ (kk_t * a[:, t])[:, :, None, :]
|
| 148 |
+
matrix = state.to(compute_dtype)
|
| 149 |
+
state = matrix * decay.to(compute_dtype)
|
| 150 |
+
state = state + matrix @ anti_value_key.to(compute_dtype)
|
| 151 |
+
state = state + value_key.to(compute_dtype)
|
| 152 |
+
if state_dtype != compute_dtype:
|
| 153 |
+
state = state.to(state_dtype)
|
| 154 |
+
out[:, t] = (state.to(dtype) @ r[:, t][..., None]).squeeze(-1).to(compute_dtype)
|
| 155 |
+
|
| 156 |
+
return out.to(dtype), state
|
| 157 |
+
|
| 158 |
+
|
| 159 |
+
def rwkv7_chunked(
|
| 160 |
+
r: torch.Tensor,
|
| 161 |
+
w_log: torch.Tensor,
|
| 162 |
+
k: torch.Tensor,
|
| 163 |
+
v: torch.Tensor,
|
| 164 |
+
kk: torch.Tensor,
|
| 165 |
+
a: torch.Tensor,
|
| 166 |
+
state: torch.Tensor,
|
| 167 |
+
chunk_size: int = 64,
|
| 168 |
+
compute_dtype: torch.dtype = torch.float32,
|
| 169 |
+
) -> tuple[torch.Tensor, torch.Tensor]:
|
| 170 |
+
"""Chunk-parallel form of [`rwkv7_recurrent`], for sequences.
|
| 171 |
+
|
| 172 |
+
The step is `S_t = A_t S_{t-1} + k_t v_t^T` with `A_t = diag(w_t) - b_t kk_t^T`
|
| 173 |
+
and `b_t = kk_t * a_t`, i.e. diagonal-plus-rank-one. Substituting
|
| 174 |
+
`S_t = diag(c_t) P_t` with the running decay `c_t = prod_{s<=t} w_s` removes the
|
| 175 |
+
diagonal part and leaves a plain delta rule::
|
| 176 |
+
|
| 177 |
+
P_t = (I - b~_t q~_t^T) P_{t-1} + k~_t v_t^T
|
| 178 |
+
b~ = b / c, k~ = k / c, q~ = kk * c_{t-1}, r~ = r * c
|
| 179 |
+
|
| 180 |
+
Writing `u_t = q~_t^T P_{t-1}` turns a whole chunk into one unit-lower-triangular
|
| 181 |
+
system, so the chunk is a handful of matmuls instead of `chunk_size` sequential
|
| 182 |
+
steps; only the chunk-to-chunk carry stays serial::
|
| 183 |
+
|
| 184 |
+
(I + tril(Q~ B~^T, -1)) U = Q~ P_0 + tril(Q~ K~^T, -1) V
|
| 185 |
+
O = R~ P_0 + tril(R~ K~^T, 0) V - tril(R~ B~^T, 0) U
|
| 186 |
+
P_C = P_0 + K~^T V - B~^T U
|
| 187 |
+
|
| 188 |
+
`chunk_size` is bounded by that division by `c`, but by OVERFLOW rather than by
|
| 189 |
+
precision, and the difference is worth a factor of four. The per-step decay is at
|
| 190 |
+
least `exp(-e^-0.5)` = 0.5452, so at the worst case, every channel pinned at that
|
| 191 |
+
floor, `1/c` grows like `e^(e^-0.5 * chunk_size)`. That reaches 7.2e16 at 64 and
|
| 192 |
+
fp32 tops out near 3.4e38, so there are twenty-odd decades of headroom left at the
|
| 193 |
+
default, and the ceiling is `ln(finfo.max) / e^-0.5` = 146 in fp32.
|
| 194 |
+
|
| 195 |
+
That derivation is worth stating carefully because the first version of it was
|
| 196 |
+
wrong in a way the numbers hid. It said the decay floor was `e^-0.5` rather than
|
| 197 |
+
`exp(-e^-0.5)`, hence growth like `e^(0.5 * chunk_size)`, which puts the ceiling at
|
| 198 |
+
177, and measured at the decay floor this function returns all-NaN from 147 up. The
|
| 199 |
+
quoted 7.2e16 came from the correct law all along (`e^(0.5*64)` is 7.9e13), so the
|
| 200 |
+
arithmetic had been done right and written up wrong, and following the prose rather
|
| 201 |
+
than the number led into the overflow band. `chunk_size` is checked below rather
|
| 202 |
+
than left to that reasoning.
|
| 203 |
+
|
| 204 |
+
Precision does not degrade with it, because the substitution is a similarity
|
| 205 |
+
transform: whatever `1/c` inflates, `c` deflates again on the way out, and a
|
| 206 |
+
common scale factor does not move a floating-point relative error. Measured at
|
| 207 |
+
the decay floor, `T=256`, against the sequential form: chunk 16 gives 2.0e-07 and
|
| 208 |
+
chunk 64 gives 2.7e-07, both fp32 noise, while the recurrence runs 3.9x faster
|
| 209 |
+
(3.771 ms -> 0.913 ms). 16 was the conservative reading of the same bound.
|
| 210 |
+
"""
|
| 211 |
+
batch, seq_len, num_heads, head_dim = r.shape
|
| 212 |
+
dtype = r.dtype
|
| 213 |
+
r, w_log, k, v, kk, a = (t.to(compute_dtype) for t in (r, w_log, k, v, kk, a))
|
| 214 |
+
state_dtype = state.dtype
|
| 215 |
+
# The chunk derivation below uses the transposed (key, value) form internally.
|
| 216 |
+
state = state.transpose(-2, -1).to(compute_dtype)
|
| 217 |
+
# Everything that does not touch the carried state is computed for ALL chunks at
|
| 218 |
+
# once. Only the state recurrence is serial, and it was pulling the rest of the
|
| 219 |
+
# arithmetic into the Python loop with it: at T=256 and chunk 64 that was four
|
| 220 |
+
# iterations of ten-odd launches each, plus three constant matrices rebuilt every
|
| 221 |
+
# iteration. The loop below now does the state-dependent terms and nothing else.
|
| 222 |
+
# Never pad past the sequence: a 16-token prefill grouped into 64-wide chunks does
|
| 223 |
+
# four times the arithmetic for the same answer, which cost 1x16 nine points before
|
| 224 |
+
# this line existed. The old loop avoided it by shortening its last chunk.
|
| 225 |
+
chunk_size = min(chunk_size, seq_len)
|
| 226 |
+
# Refuse rather than return NaN. `1/c` reaches `e^(e^-0.5 * chunk_size)` at the
|
| 227 |
+
# decay floor, so the widest chunk this dtype can carry is derived from the dtype
|
| 228 |
+
# rather than written down: 146 in fp32, 12 in fp16. A caller who raises
|
| 229 |
+
# `chunk_size` for a longer prefill gets an error naming the limit instead of
|
| 230 |
+
# all-NaN logits several layers later, which is what the previous version did.
|
| 231 |
+
widest = int(math.log(torch.finfo(compute_dtype).max) / _INV_SQRT_E)
|
| 232 |
+
if chunk_size > widest:
|
| 233 |
+
raise ValueError(
|
| 234 |
+
f"chunk_size={chunk_size} overflows {compute_dtype}: the running decay's "
|
| 235 |
+
f"reciprocal grows like e^(e^-0.5 * chunk_size), so the widest chunk that "
|
| 236 |
+
f"stays finite is {widest}. Lower chunk_size, or raise compute_dtype."
|
| 237 |
+
)
|
| 238 |
+
chunks = (seq_len + chunk_size - 1) // chunk_size
|
| 239 |
+
padded = chunks * chunk_size
|
| 240 |
+
if padded != seq_len:
|
| 241 |
+
pad = (0, 0, 0, 0, 0, padded - seq_len)
|
| 242 |
+
r, k, v, kk, a = (torch.nn.functional.pad(t, pad) for t in (r, k, v, kk, a))
|
| 243 |
+
# A pad step must be the identity for the recurrence: decay 1, nothing added.
|
| 244 |
+
w_log = torch.nn.functional.pad(w_log, pad)
|
| 245 |
+
grouped = lambda t: t.reshape(batch, chunks, chunk_size, num_heads, head_dim) # noqa: E731
|
| 246 |
+
rg, kg, vg, kkg, ag, wg = (grouped(t) for t in (r, k, v, kk, a, w_log))
|
| 247 |
+
|
| 248 |
+
w_c = torch.exp(wg)
|
| 249 |
+
c = torch.cumprod(w_c, dim=2)
|
| 250 |
+
c_prev = c / w_c
|
| 251 |
+
bg = kkg * ag
|
| 252 |
+
b_t, k_t = bg / c, kg / c
|
| 253 |
+
q_t, r_t = kkg * c_prev, rg * c
|
| 254 |
+
|
| 255 |
+
span = chunk_size
|
| 256 |
+
tri = torch.ones(span, span, device=r.device, dtype=r.dtype).tril(-1)
|
| 257 |
+
causal = torch.ones(span, span, device=r.device, dtype=r.dtype).tril(0)
|
| 258 |
+
eye = torch.eye(span, device=r.device, dtype=r.dtype)
|
| 259 |
+
|
| 260 |
+
# [batch, chunks, heads, span, span] -- one launch each instead of one per chunk.
|
| 261 |
+
qb = torch.einsum("bcthn,bcshn->bchts", q_t, b_t) * tri
|
| 262 |
+
qk = torch.einsum("bcthn,bcshn->bchts", q_t, k_t) * tri
|
| 263 |
+
rk = torch.einsum("bcthn,bcshn->bchts", r_t, k_t) * causal
|
| 264 |
+
rb = torch.einsum("bcthn,bcshn->bchts", r_t, b_t) * causal
|
| 265 |
+
lhs = eye + qb
|
| 266 |
+
qkv = torch.einsum("bchts,bcshv->bchtv", qk, vg)
|
| 267 |
+
rkv = torch.einsum("bchts,bcshv->bchtv", rk, vg)
|
| 268 |
+
c_last = c[:, :, -1].unsqueeze(-1)
|
| 269 |
+
|
| 270 |
+
# `k~^T v` carries no state either, so it joins the batched half above; only the
|
| 271 |
+
# two products that read the carried state and the solve that depends on them are
|
| 272 |
+
# genuinely serial.
|
| 273 |
+
kv = torch.einsum("bcthn,bcthv->bchnv", k_t, vg)
|
| 274 |
+
# Both state products read the SAME state, so they are one matmul over a stacked
|
| 275 |
+
# token axis rather than two -- the loop is short and launch-bound, and this is one
|
| 276 |
+
# launch per chunk instead of two.
|
| 277 |
+
qr = torch.cat([q_t, r_t], dim=2)
|
| 278 |
+
|
| 279 |
+
outputs = []
|
| 280 |
+
for i in range(chunks):
|
| 281 |
+
qr_s = torch.einsum("bthn,bhnv->bhtv", qr[:, i], state)
|
| 282 |
+
rhs = qr_s[:, :, :chunk_size] + qkv[:, i]
|
| 283 |
+
u = torch.linalg.solve_triangular(lhs[:, i], rhs, upper=False, unitriangular=True)
|
| 284 |
+
out_c = qr_s[:, :, chunk_size:] + rkv[:, i] - torch.einsum("bhts,bhsv->bhtv", rb[:, i], u)
|
| 285 |
+
outputs.append(out_c.permute(0, 2, 1, 3))
|
| 286 |
+
state = c_last[:, i] * (state + kv[:, i] - torch.einsum("bthn,bhtv->bhnv", b_t[:, i], u))
|
| 287 |
+
|
| 288 |
+
state = state.transpose(-2, -1).contiguous().to(state_dtype)
|
| 289 |
+
return torch.cat(outputs, dim=1)[:, :seq_len].to(dtype), state
|
| 290 |
+
|
| 291 |
+
|
| 292 |
+
def rwkv7_eager(
|
| 293 |
+
r: torch.Tensor,
|
| 294 |
+
w_log: torch.Tensor,
|
| 295 |
+
k: torch.Tensor,
|
| 296 |
+
v: torch.Tensor,
|
| 297 |
+
kk: torch.Tensor,
|
| 298 |
+
a: torch.Tensor,
|
| 299 |
+
state: torch.Tensor,
|
| 300 |
+
cu_seq_lens: torch.Tensor | None = None,
|
| 301 |
+
**kwargs,
|
| 302 |
+
) -> tuple[torch.Tensor, torch.Tensor]:
|
| 303 |
+
"""Portable exact WKV, evaluated in reference token order.
|
| 304 |
+
|
| 305 |
+
With `cu_seq_lens` the row is a *packed* batch: several independent sequences
|
| 306 |
+
laid end to end, the layout a varlen kernel consumes. Each one has to start
|
| 307 |
+
from a fresh state, so they are run in turn and their outputs concatenated;
|
| 308 |
+
the state returned is the last segment's, which is the one a continuation
|
| 309 |
+
would resume from. This is the reference behaviour, not the fast path: a
|
| 310 |
+
fused varlen kernel does the same segments in one launch, which is exactly
|
| 311 |
+
the kind of implementation [`RWKV7_WKV_FUNCTIONS`] exists to let you drop in.
|
| 312 |
+
|
| 313 |
+
This default deliberately stays sequential for a multi-token prefill. The
|
| 314 |
+
chunk-parallel rearrangement is much faster, but changes floating-point operation
|
| 315 |
+
order and is therefore available only through the explicit `"chunked"` registry
|
| 316 |
+
entry. Automatic execution must remain the reference users can compare it to.
|
| 317 |
+
"""
|
| 318 |
+
if cu_seq_lens is not None:
|
| 319 |
+
bounds = cu_seq_lens.tolist()
|
| 320 |
+
outputs = []
|
| 321 |
+
for start, stop in zip(bounds[:-1], bounds[1:]):
|
| 322 |
+
if stop <= start:
|
| 323 |
+
continue
|
| 324 |
+
segment_state = torch.zeros_like(state)
|
| 325 |
+
out, state = rwkv7_eager(
|
| 326 |
+
r[:, start:stop],
|
| 327 |
+
w_log[:, start:stop],
|
| 328 |
+
k[:, start:stop],
|
| 329 |
+
v[:, start:stop],
|
| 330 |
+
kk[:, start:stop],
|
| 331 |
+
a[:, start:stop],
|
| 332 |
+
segment_state,
|
| 333 |
+
)
|
| 334 |
+
outputs.append(out)
|
| 335 |
+
return torch.cat(outputs, dim=1), state
|
| 336 |
+
|
| 337 |
+
return rwkv7_recurrent(r, w_log, k, v, kk, a, state)
|
| 338 |
+
|
| 339 |
+
|
| 340 |
+
def rwkv7_parallel(
|
| 341 |
+
r: torch.Tensor,
|
| 342 |
+
w_log: torch.Tensor,
|
| 343 |
+
k: torch.Tensor,
|
| 344 |
+
v: torch.Tensor,
|
| 345 |
+
kk: torch.Tensor,
|
| 346 |
+
a: torch.Tensor,
|
| 347 |
+
state: torch.Tensor,
|
| 348 |
+
cu_seq_lens: torch.Tensor | None = None,
|
| 349 |
+
**kwargs,
|
| 350 |
+
) -> tuple[torch.Tensor, torch.Tensor]:
|
| 351 |
+
"""Explicit chunk-parallel WKV with the same packed-sequence contract as eager."""
|
| 352 |
+
if cu_seq_lens is not None:
|
| 353 |
+
bounds = cu_seq_lens.tolist()
|
| 354 |
+
outputs = []
|
| 355 |
+
for start, stop in zip(bounds[:-1], bounds[1:]):
|
| 356 |
+
if stop <= start:
|
| 357 |
+
continue
|
| 358 |
+
out, state = rwkv7_chunked(
|
| 359 |
+
r[:, start:stop],
|
| 360 |
+
w_log[:, start:stop],
|
| 361 |
+
k[:, start:stop],
|
| 362 |
+
v[:, start:stop],
|
| 363 |
+
kk[:, start:stop],
|
| 364 |
+
a[:, start:stop],
|
| 365 |
+
torch.zeros_like(state),
|
| 366 |
+
)
|
| 367 |
+
outputs.append(out)
|
| 368 |
+
return torch.cat(outputs, dim=1), state
|
| 369 |
+
return rwkv7_chunked(r, w_log, k, v, kk, a, state)
|
| 370 |
+
|
| 371 |
+
|
| 372 |
+
RWKV7_WKV_FUNCTIONS = {"eager": rwkv7_eager, "chunked": rwkv7_parallel}
|
| 373 |
+
|
| 374 |
+
|
| 375 |
+
class Rwkv7Attention(nn.Module):
|
| 376 |
+
"""RWKV-7 time-mixing block (the recurrent replacement for self-attention)."""
|
| 377 |
+
|
| 378 |
+
def __init__(self, config: Rwkv7Config, layer_id: int):
|
| 379 |
+
super().__init__()
|
| 380 |
+
self.config = config
|
| 381 |
+
self.layer_id = layer_id
|
| 382 |
+
C = config.hidden_size
|
| 383 |
+
self.hidden_size = C
|
| 384 |
+
self.num_heads = config.num_heads
|
| 385 |
+
self.head_dim = config.head_dim
|
| 386 |
+
self.state_dtype = getattr(torch, config.wkv_state_dtype)
|
| 387 |
+
|
| 388 |
+
self.time_shift = Rwkv7TokenShift()
|
| 389 |
+
|
| 390 |
+
# per-channel token-shift mixes (kept at Bo's (1, 1, C) shape)
|
| 391 |
+
self.x_r = nn.Parameter(torch.zeros(1, 1, C))
|
| 392 |
+
self.x_w = nn.Parameter(torch.zeros(1, 1, C))
|
| 393 |
+
self.x_k = nn.Parameter(torch.zeros(1, 1, C))
|
| 394 |
+
self.x_v = nn.Parameter(torch.zeros(1, 1, C))
|
| 395 |
+
self.x_a = nn.Parameter(torch.zeros(1, 1, C))
|
| 396 |
+
self.x_g = nn.Parameter(torch.zeros(1, 1, C))
|
| 397 |
+
|
| 398 |
+
# LoRA factors as raw tensors, exactly as the reference stores them
|
| 399 |
+
self.w1 = nn.Parameter(torch.zeros(C, config.decay_low_rank_dim))
|
| 400 |
+
self.w2 = nn.Parameter(torch.zeros(config.decay_low_rank_dim, C))
|
| 401 |
+
self.w0 = nn.Parameter(torch.zeros(1, 1, C))
|
| 402 |
+
self.a1 = nn.Parameter(torch.zeros(C, config.a_low_rank_dim))
|
| 403 |
+
self.a2 = nn.Parameter(torch.zeros(config.a_low_rank_dim, C))
|
| 404 |
+
self.a0 = nn.Parameter(torch.zeros(1, 1, C))
|
| 405 |
+
self.g1 = nn.Parameter(torch.zeros(C, config.gate_low_rank_dim))
|
| 406 |
+
self.g2 = nn.Parameter(torch.zeros(config.gate_low_rank_dim, C))
|
| 407 |
+
# The value-residual LoRA exists on every layer in a reference checkpoint,
|
| 408 |
+
# but layer 0 PRODUCES v_first instead of mixing towards it, so its copy is
|
| 409 |
+
# never read. It is registered anyway so that loading is lossless.
|
| 410 |
+
self.v1 = nn.Parameter(torch.zeros(C, config.v_low_rank_dim))
|
| 411 |
+
self.v2 = nn.Parameter(torch.zeros(config.v_low_rank_dim, C))
|
| 412 |
+
self.v0 = nn.Parameter(torch.zeros(1, 1, C))
|
| 413 |
+
|
| 414 |
+
self.k_k = nn.Parameter(torch.zeros(1, 1, C))
|
| 415 |
+
self.k_a = nn.Parameter(torch.zeros(1, 1, C))
|
| 416 |
+
self.r_k = nn.Parameter(torch.zeros(config.num_heads, config.head_dim))
|
| 417 |
+
|
| 418 |
+
self.receptance = nn.Linear(C, C, bias=False)
|
| 419 |
+
self.key = nn.Linear(C, C, bias=False)
|
| 420 |
+
self.value = nn.Linear(C, C, bias=False)
|
| 421 |
+
self.output = nn.Linear(C, C, bias=False)
|
| 422 |
+
# GroupNorm over heads, matching the reference's per-head normalisation.
|
| 423 |
+
#
|
| 424 |
+
# The reference hardcodes `eps=64e-5`, which is `head_dim * 1e-5` at its
|
| 425 |
+
# head_dim of 64 -- NOT `num_heads * 1e-5`. The two coincide only when a
|
| 426 |
+
# model happens to have as many heads as channels per head, which for
|
| 427 |
+
# head_dim 64 means hidden_size 4096 exactly. Scaling by `num_heads` is
|
| 428 |
+
# therefore right on the 7.2B and wrong everywhere else: 2x low on the 1.5B,
|
| 429 |
+
# 5.33x low on the 0.1B. It is written against `head_dim` because that is
|
| 430 |
+
# the axis GroupNorm actually reduces over, so the constant tracks the
|
| 431 |
+
# reference at any width.
|
| 432 |
+
self.ln_x = nn.GroupNorm(config.num_heads, C, eps=config.norm_eps * config.head_dim)
|
| 433 |
+
|
| 434 |
+
def forward(
|
| 435 |
+
self,
|
| 436 |
+
hidden_states: torch.Tensor,
|
| 437 |
+
v_first: torch.Tensor | None,
|
| 438 |
+
shift_state: torch.Tensor | None,
|
| 439 |
+
wkv_state: torch.Tensor | None,
|
| 440 |
+
keep: torch.Tensor | None = None,
|
| 441 |
+
cu_seq_lens: torch.Tensor | None = None,
|
| 442 |
+
):
|
| 443 |
+
batch, seq_len, C = hidden_states.shape
|
| 444 |
+
H, N = self.num_heads, self.head_dim
|
| 445 |
+
tokenwise = self.config.wkv_implementation != "chunked"
|
| 446 |
+
|
| 447 |
+
shifted, new_shift_state = self.time_shift(hidden_states, shift_state, cu_seq_lens, keep)
|
| 448 |
+
delta = shifted - hidden_states
|
| 449 |
+
xr = hidden_states + self.x_r * delta
|
| 450 |
+
xw = hidden_states + self.x_w * delta
|
| 451 |
+
xk = hidden_states + self.x_k * delta
|
| 452 |
+
xv = hidden_states + self.x_v * delta
|
| 453 |
+
xa = hidden_states + self.x_a * delta
|
| 454 |
+
xg = hidden_states + self.x_g * delta
|
| 455 |
+
|
| 456 |
+
r = _rwkv7_linear(self.receptance, xr, tokenwise)
|
| 457 |
+
k = _rwkv7_linear(self.key, xk, tokenwise)
|
| 458 |
+
v = _rwkv7_linear(self.value, xv, tokenwise)
|
| 459 |
+
|
| 460 |
+
w_log, a, g, v_gate = self.lora_gates(xw, xa, xg, None if self.layer_id == 0 else xv, tokenwise)
|
| 461 |
+
|
| 462 |
+
if keep is not None:
|
| 463 |
+
# A padding position has to leave the recurrent state exactly as it
|
| 464 |
+
# found it, which takes three things and not one.
|
| 465 |
+
#
|
| 466 |
+
# The decay is held at w = exp(0) = 1, so the transition is the identity.
|
| 467 |
+
#
|
| 468 |
+
# `k` and `v` are zeroed EXPLICITLY rather than relying on the blanked
|
| 469 |
+
# hidden state to make them zero: the projections are bias-free, so that
|
| 470 |
+
# only holds for a pad with nothing before it.
|
| 471 |
+
# A pad that FOLLOWS a real token still receives that token's hidden
|
| 472 |
+
# state through the token shift, so `delta = shifted - 0` is non-zero,
|
| 473 |
+
# and so are `k` and `v`. The update term `k v^T` then entered the state
|
| 474 |
+
# on every right-padded batch and on any left-padded batch continued
|
| 475 |
+
# from a carried shift state.
|
| 476 |
+
w_log = w_log * keep
|
| 477 |
+
k = k * keep
|
| 478 |
+
v = v * keep
|
| 479 |
+
|
| 480 |
+
if self.layer_id == 0:
|
| 481 |
+
v_first = v
|
| 482 |
+
else:
|
| 483 |
+
v = v + (v_first - v) * v_gate
|
| 484 |
+
|
| 485 |
+
kk = k * self.k_k
|
| 486 |
+
kk = torch.nn.functional.normalize(kk.view(batch, seq_len, H, N), dim=-1, p=2.0).view(batch, seq_len, C)
|
| 487 |
+
if keep is not None:
|
| 488 |
+
# A blanked padding position makes `k` exactly zero, so this normalises a
|
| 489 |
+
# zero vector. `F.normalize` divides by `max(norm, 1e-12)`, and 1e-12 is
|
| 490 |
+
# below the smallest fp16 subnormal -- so in fp16 the divisor really is
|
| 491 |
+
# zero and every padded row comes out NaN. (In fp32 it is representable
|
| 492 |
+
# and the whole thing looks fine, which is why a test on a small fp32
|
| 493 |
+
# model missed it entirely.) `where` is used rather than a multiply
|
| 494 |
+
# because NaN * 0 is still NaN.
|
| 495 |
+
kk = torch.where(keep.bool().expand_as(kk), kk, torch.zeros_like(kk))
|
| 496 |
+
k = k * (1 + (a - 1) * self.k_a)
|
| 497 |
+
|
| 498 |
+
if wkv_state is None:
|
| 499 |
+
wkv_state = torch.zeros(batch, H, N, N, device=r.device, dtype=self.state_dtype)
|
| 500 |
+
|
| 501 |
+
def _heads(t):
|
| 502 |
+
return t.view(batch, seq_len, H, N)
|
| 503 |
+
|
| 504 |
+
# Named rather than indexed: an unregistered key would otherwise surface as a
|
| 505 |
+
# bare `KeyError: 'chunked'` from inside the forward, several frames from the
|
| 506 |
+
# config field that caused it, and a caller looping over shapes would record it
|
| 507 |
+
# as "this shape did not run" instead of "this model was never built". The
|
| 508 |
+
# registry is open by design, so this cannot be validated in the config.
|
| 509 |
+
try:
|
| 510 |
+
wkv = RWKV7_WKV_FUNCTIONS[self.config.wkv_implementation]
|
| 511 |
+
except KeyError:
|
| 512 |
+
raise ValueError(
|
| 513 |
+
f"wkv_implementation={self.config.wkv_implementation!r} is not registered. "
|
| 514 |
+
f"Known: {sorted(RWKV7_WKV_FUNCTIONS)}. Register your own with "
|
| 515 |
+
"`RWKV7_WKV_FUNCTIONS['name'] = fn` before building the model."
|
| 516 |
+
) from None
|
| 517 |
+
y, wkv_state = wkv(
|
| 518 |
+
_heads(r),
|
| 519 |
+
_heads(w_log),
|
| 520 |
+
_heads(k),
|
| 521 |
+
_heads(v),
|
| 522 |
+
_heads(kk),
|
| 523 |
+
_heads(a),
|
| 524 |
+
wkv_state,
|
| 525 |
+
cu_seq_lens=cu_seq_lens,
|
| 526 |
+
)
|
| 527 |
+
|
| 528 |
+
# `reshape`, not `view`: what comes back from the WKV is whatever layout that
|
| 529 |
+
# implementation produced, and a registered kernel -- or inductor, which is
|
| 530 |
+
# free to pick its own -- can hand back a `[batch, seq_len, heads, head_dim]`
|
| 531 |
+
# that is strided as if it were `[batch, heads, seq_len, head_dim]`. `view`
|
| 532 |
+
# then raises, and it raises only when batch and seq_len are BOTH greater
|
| 533 |
+
# than one, which is why a compiled 16x16 forward failed while 1xT and Bx1
|
| 534 |
+
# both passed.
|
| 535 |
+
y = self.ln_x(y.reshape(batch * seq_len, C)).view(batch, seq_len, C)
|
| 536 |
+
# r·k·r_k summed per head, broadcast back over the head's value channels
|
| 537 |
+
bonus = ((_heads(r) * _heads(k) * self.r_k).sum(dim=-1, keepdim=True) * _heads(v)).reshape(batch, seq_len, C)
|
| 538 |
+
y = _rwkv7_linear(self.output, (y + bonus) * g, tokenwise)
|
| 539 |
+
return y, v_first, new_shift_state, wkv_state
|
| 540 |
+
|
| 541 |
+
def lora_gates(self, xw, xa, xg, xv, tokenwise):
|
| 542 |
+
"""The four low-rank chains, as one unit.
|
| 543 |
+
|
| 544 |
+
Each is a rank-r down projection, an activation, and an up projection, and
|
| 545 |
+
they differ only in where the activation sits: `w` has `tanh` between the
|
| 546 |
+
two, `g` has `sigmoid` between them and nothing after, `a` and `v` have
|
| 547 |
+
nothing between and `sigmoid` after. Together they are a dozen tiny
|
| 548 |
+
matrix-vector products per layer whose weights are a rounding error of the
|
| 549 |
+
model's bytes -- so at batch 1 they cost latency, not bandwidth, which is
|
| 550 |
+
the shape a fused kernel improves and a portable implementation cannot.
|
| 551 |
+
|
| 552 |
+
Kept as one method for exactly that reason: it is the unit worth replacing,
|
| 553 |
+
and replacing it should not need a fork. `xv` is None on layer 0, which
|
| 554 |
+
produces `v_first` rather than mixing towards it and never reads that chain.
|
| 555 |
+
"""
|
| 556 |
+
w_log = -_INV_SQRT_E * torch.sigmoid(
|
| 557 |
+
_rwkv7_matmul(torch.tanh(_rwkv7_matmul(xw, self.w1, tokenwise)), self.w2, tokenwise) + self.w0
|
| 558 |
+
)
|
| 559 |
+
a = torch.sigmoid(_rwkv7_matmul(_rwkv7_matmul(xa, self.a1, tokenwise), self.a2, tokenwise) + self.a0)
|
| 560 |
+
g = _rwkv7_matmul(torch.sigmoid(_rwkv7_matmul(xg, self.g1, tokenwise)), self.g2, tokenwise)
|
| 561 |
+
v_gate = (
|
| 562 |
+
None
|
| 563 |
+
if xv is None
|
| 564 |
+
else torch.sigmoid(_rwkv7_matmul(_rwkv7_matmul(xv, self.v1, tokenwise), self.v2, tokenwise) + self.v0)
|
| 565 |
+
)
|
| 566 |
+
return w_log, a, g, v_gate
|
| 567 |
+
|
| 568 |
+
|
| 569 |
+
class Rwkv7FeedForward(nn.Module):
|
| 570 |
+
"""RWKV-7 channel-mixing block: squared-ReLU over a single token shift."""
|
| 571 |
+
|
| 572 |
+
def __init__(self, config: Rwkv7Config, layer_id: int):
|
| 573 |
+
super().__init__()
|
| 574 |
+
self.config = config
|
| 575 |
+
self.layer_id = layer_id
|
| 576 |
+
C = config.hidden_size
|
| 577 |
+
self.time_shift = Rwkv7TokenShift()
|
| 578 |
+
self.x_k = nn.Parameter(torch.zeros(1, 1, C))
|
| 579 |
+
self.key = nn.Linear(C, config.intermediate_size, bias=False)
|
| 580 |
+
self.value = nn.Linear(config.intermediate_size, C, bias=False)
|
| 581 |
+
|
| 582 |
+
def forward(
|
| 583 |
+
self,
|
| 584 |
+
hidden_states: torch.Tensor,
|
| 585 |
+
shift_state: torch.Tensor | None,
|
| 586 |
+
cu_seq_lens: torch.Tensor | None = None,
|
| 587 |
+
keep: torch.Tensor | None = None,
|
| 588 |
+
):
|
| 589 |
+
tokenwise = self.config.wkv_implementation != "chunked"
|
| 590 |
+
shifted, new_shift_state = self.time_shift(hidden_states, shift_state, cu_seq_lens, keep)
|
| 591 |
+
xk = hidden_states + self.x_k * (shifted - hidden_states)
|
| 592 |
+
inner = torch.relu(_rwkv7_linear(self.key, xk, tokenwise)) ** 2
|
| 593 |
+
out = self._project(inner, tokenwise)
|
| 594 |
+
return out, new_shift_state
|
| 595 |
+
|
| 596 |
+
def _project(self, inner: torch.Tensor, tokenwise: bool) -> torch.Tensor:
|
| 597 |
+
return _rwkv7_linear(self.value, inner, tokenwise)
|
| 598 |
+
|
| 599 |
+
|
| 600 |
+
class Rwkv7CacheLayer(LinearAttentionLayer):
|
| 601 |
+
"""One block's slice of the recurrent state: the WKV matrix and two token shifts.
|
| 602 |
+
|
| 603 |
+
Everything a beam search or a batched generate needs to do to a cache is a
|
| 604 |
+
permutation of its batch axis, and for RWKV-7 that is the whole job -- the state
|
| 605 |
+
is O(1) in sequence length, so there is no time axis to gather along and no
|
| 606 |
+
length bookkeeping to keep consistent.
|
| 607 |
+
|
| 608 |
+
The slot layout is [`Rwkv7Cache`]'s; this class only moves whatever is in them.
|
| 609 |
+
"""
|
| 610 |
+
|
| 611 |
+
def lazy_initialization(self, conv_states=None, recurrent_states=None, state_idx: int = 0, **kwargs) -> None:
|
| 612 |
+
super().lazy_initialization(conv_states, recurrent_states, state_idx, **kwargs)
|
| 613 |
+
# Upstream records device/dtype only on the conv branch, since a linear
|
| 614 |
+
# attention layer normally has a convolution in front of it. This model has
|
| 615 |
+
# none, so without this every recurrent-only layer would keep `device=None`
|
| 616 |
+
# and `reorder_cache`'s `beam_idx.to(self.device)` would fail.
|
| 617 |
+
if recurrent_states is not None and self.device is None:
|
| 618 |
+
self.dtype, self.device = recurrent_states.dtype, recurrent_states.device
|
| 619 |
+
|
| 620 |
+
def allocate(
|
| 621 |
+
self, batch: int, shapes: dict[int, tuple], device, dtypes: dict[int, torch.dtype], state_dtype
|
| 622 |
+
) -> None:
|
| 623 |
+
"""Create every slot up front, zeroed, at pinned addresses.
|
| 624 |
+
|
| 625 |
+
The lazy path allocates a slot the first time it is written, which is inside
|
| 626 |
+
the compiled region -- and a buffer that first appears there cannot be given
|
| 627 |
+
a static address, so inductor declines CUDA graphs for a recurrent decode,
|
| 628 |
+
the one workload that most needs them.
|
| 629 |
+
"""
|
| 630 |
+
for slot, shape in shapes.items():
|
| 631 |
+
buffer = torch.zeros((batch, *shape), device=device, dtype=dtypes[slot])
|
| 632 |
+
if not torch.compiler.is_compiling():
|
| 633 |
+
torch._dynamo.mark_static_address(buffer)
|
| 634 |
+
self.recurrent_states[slot] = buffer
|
| 635 |
+
self.is_recurrent_states_initialized[slot] = True
|
| 636 |
+
self.has_previous_state[slot] = True
|
| 637 |
+
self.dtype, self.device = state_dtype, device
|
| 638 |
+
|
| 639 |
+
def reorder_cache(self, beam_idx: torch.LongTensor) -> None:
|
| 640 |
+
"""Permute the batch onto the beams that survived, without moving house.
|
| 641 |
+
|
| 642 |
+
Upstream rebinds each slot to the `index_select` result. That is correct but
|
| 643 |
+
it hands back a freshly allocated tensor, so the address pinned at allocation
|
| 644 |
+
is gone for the rest of the generation and the compiled decode quietly loses
|
| 645 |
+
its CUDA graphs. Copying back into the same buffer costs one temporary and
|
| 646 |
+
keeps the pinning.
|
| 647 |
+
"""
|
| 648 |
+
for slot in range(self.number_of_states):
|
| 649 |
+
buffer = self.recurrent_states[slot]
|
| 650 |
+
if self.is_recurrent_states_initialized[slot] and buffer is not None:
|
| 651 |
+
buffer.copy_(buffer.index_select(0, beam_idx.to(buffer.device)))
|
| 652 |
+
|
| 653 |
+
def batch_repeat_interleave(self, repeats: int) -> None:
|
| 654 |
+
"""Fan each sequence's state out to `repeats` copies, for one-prompt-many-samples.
|
| 655 |
+
|
| 656 |
+
This changes the batch size, so unlike `reorder_cache` it cannot preserve the
|
| 657 |
+
pinned addresses -- the buffers are necessarily new ones. Callers that then
|
| 658 |
+
compile should re-pin, which `Rwkv7Model.allocate_state` does.
|
| 659 |
+
"""
|
| 660 |
+
for slot in range(self.number_of_states):
|
| 661 |
+
buffer = self.recurrent_states[slot]
|
| 662 |
+
if self.is_recurrent_states_initialized[slot] and buffer is not None:
|
| 663 |
+
self.recurrent_states[slot] = buffer.repeat_interleave(repeats, dim=0)
|
| 664 |
+
|
| 665 |
+
def batch_select_indices(self, indices: torch.Tensor) -> None:
|
| 666 |
+
"""Keep only the given batch rows. Same address caveat as `batch_repeat_interleave`."""
|
| 667 |
+
for slot in range(self.number_of_states):
|
| 668 |
+
buffer = self.recurrent_states[slot]
|
| 669 |
+
if self.is_recurrent_states_initialized[slot] and buffer is not None:
|
| 670 |
+
self.recurrent_states[slot] = buffer[indices, ...]
|
| 671 |
+
|
| 672 |
+
|
| 673 |
+
class Rwkv7Cache(Cache):
|
| 674 |
+
"""The recurrent state of every block, as a `Cache`.
|
| 675 |
+
|
| 676 |
+
It replaces a KV cache and is a constant size: one `[num_heads, head_dim,
|
| 677 |
+
head_dim]` matrix and two `[hidden]` token shifts per layer per sequence,
|
| 678 |
+
whatever the context length. That is the property the architecture is for, so
|
| 679 |
+
`get_max_length()` is -1 (no limit) and nothing here grows as tokens arrive.
|
| 680 |
+
"""
|
| 681 |
+
|
| 682 |
+
# Slot numbering inside one block's layer. All three are recurrent states:
|
| 683 |
+
# RWKV-7's token shift is a one-token history, not a convolution window, so it
|
| 684 |
+
# lives in a recurrent slot, where the update is a plain copy rather than the
|
| 685 |
+
# rolling concatenate a conv slot would do.
|
| 686 |
+
WKV, ATT_SHIFT, FFN_SHIFT = 0, 1, 2
|
| 687 |
+
|
| 688 |
+
def __init__(
|
| 689 |
+
self,
|
| 690 |
+
config: Rwkv7Config,
|
| 691 |
+
batch_size: int | None = None,
|
| 692 |
+
device: torch.device | str | None = None,
|
| 693 |
+
dtype: torch.dtype | None = None,
|
| 694 |
+
):
|
| 695 |
+
super().__init__(layers=[Rwkv7CacheLayer(number_of_states=3) for _ in range(config.num_hidden_layers)])
|
| 696 |
+
self.config = config
|
| 697 |
+
if batch_size is not None:
|
| 698 |
+
self.allocate(batch_size, device, dtype)
|
| 699 |
+
|
| 700 |
+
def allocate(self, batch_size: int, device=None, dtype=None) -> "Rwkv7Cache":
|
| 701 |
+
config = self.config
|
| 702 |
+
shapes = {
|
| 703 |
+
self.WKV: (config.num_heads, config.head_dim, config.head_dim),
|
| 704 |
+
self.ATT_SHIFT: (config.hidden_size,),
|
| 705 |
+
self.FFN_SHIFT: (config.hidden_size,),
|
| 706 |
+
}
|
| 707 |
+
# The WKV state carries the whole history and is the one place where a
|
| 708 |
+
# narrower dtype actually costs accuracy, so it is configured separately
|
| 709 |
+
# from the activation dtype the shifts follow.
|
| 710 |
+
dtypes = {
|
| 711 |
+
self.WKV: getattr(torch, config.wkv_state_dtype),
|
| 712 |
+
self.ATT_SHIFT: dtype,
|
| 713 |
+
self.FFN_SHIFT: dtype,
|
| 714 |
+
}
|
| 715 |
+
for layer in self.layers:
|
| 716 |
+
assert isinstance(layer, Rwkv7CacheLayer)
|
| 717 |
+
layer.allocate(batch_size, shapes, device, dtypes, dtypes[self.WKV])
|
| 718 |
+
return self
|
| 719 |
+
|
| 720 |
+
def read(self, layer_idx: int):
|
| 721 |
+
"""This block's `(att_shift, ffn_shift, wkv)`, each None before allocation."""
|
| 722 |
+
layer = self.layers[layer_idx]
|
| 723 |
+
assert isinstance(layer, Rwkv7CacheLayer)
|
| 724 |
+
states = layer.recurrent_states
|
| 725 |
+
return states[self.ATT_SHIFT], states[self.FFN_SHIFT], states[self.WKV]
|
| 726 |
+
|
| 727 |
+
def write(self, layer_idx: int, att_shift: torch.Tensor, ffn_shift: torch.Tensor, wkv: torch.Tensor) -> None:
|
| 728 |
+
states = (
|
| 729 |
+
(self.ATT_SHIFT, att_shift),
|
| 730 |
+
(self.FFN_SHIFT, ffn_shift),
|
| 731 |
+
(self.WKV, wkv),
|
| 732 |
+
)
|
| 733 |
+
# Copying into the pre-allocated slot is what keeps its address fixed, which
|
| 734 |
+
# is what lets a captured CUDA graph replay the decode loop. It is also an
|
| 735 |
+
# in-place write on a tensor autograd may be holding, and backward through a
|
| 736 |
+
# cached forward then dies on "a variable needed for gradient computation has
|
| 737 |
+
# been modified by an inplace operation". Rebinding instead costs nothing
|
| 738 |
+
# here, because a training step is not the workload that wants graphs.
|
| 739 |
+
if torch.is_grad_enabled() and any(state.requires_grad for _, state in states):
|
| 740 |
+
layer = self.layers[layer_idx]
|
| 741 |
+
assert isinstance(layer, Rwkv7CacheLayer)
|
| 742 |
+
for slot, state in states:
|
| 743 |
+
layer.recurrent_states[slot] = state
|
| 744 |
+
layer.is_recurrent_states_initialized[slot] = True
|
| 745 |
+
layer.has_previous_state[slot] = True
|
| 746 |
+
return
|
| 747 |
+
for slot, state in states:
|
| 748 |
+
self.update_recurrent_state(state, layer_idx, slot)
|
| 749 |
+
|
| 750 |
+
|
| 751 |
+
class Rwkv7Block(GradientCheckpointingLayer):
|
| 752 |
+
def __init__(self, config: Rwkv7Config, layer_id: int):
|
| 753 |
+
super().__init__()
|
| 754 |
+
self.config = config
|
| 755 |
+
self.layer_id = layer_id
|
| 756 |
+
eps, bias = config.norm_eps, config.norm_bias
|
| 757 |
+
# Layer 0 carries the extra input norm of the reference implementation.
|
| 758 |
+
if layer_id == 0:
|
| 759 |
+
self.ln0 = nn.LayerNorm(config.hidden_size, eps=eps, bias=bias)
|
| 760 |
+
self.ln1 = nn.LayerNorm(config.hidden_size, eps=eps, bias=bias)
|
| 761 |
+
self.ln2 = nn.LayerNorm(config.hidden_size, eps=eps, bias=bias)
|
| 762 |
+
self.att = Rwkv7Attention(config, layer_id)
|
| 763 |
+
self.ffn = Rwkv7FeedForward(config, layer_id)
|
| 764 |
+
|
| 765 |
+
def forward(
|
| 766 |
+
self,
|
| 767 |
+
hidden_states: torch.Tensor,
|
| 768 |
+
v_first: torch.Tensor | None,
|
| 769 |
+
state: Rwkv7Cache | None,
|
| 770 |
+
keep: torch.Tensor | None = None,
|
| 771 |
+
cu_seq_lens: torch.Tensor | None = None,
|
| 772 |
+
):
|
| 773 |
+
if self.layer_id == 0:
|
| 774 |
+
hidden_states = self.ln0(hidden_states)
|
| 775 |
+
|
| 776 |
+
att_shift, ffn_shift, wkv = state.read(self.layer_id) if state is not None else (None, None, None)
|
| 777 |
+
|
| 778 |
+
# Padding is blanked AFTER each norm, not before: a LayerNorm maps the zero
|
| 779 |
+
# vector to its own bias, so masking the residual stream instead would let
|
| 780 |
+
# every pad position come back to life on the way into the next mixer -- and
|
| 781 |
+
# a live pad both moves the state and leaks into the next token's shift.
|
| 782 |
+
attn_in = self.ln1(hidden_states)
|
| 783 |
+
if keep is not None:
|
| 784 |
+
attn_in = attn_in * keep
|
| 785 |
+
attn_out, v_first, att_shift, wkv = self.att(attn_in, v_first, att_shift, wkv, keep, cu_seq_lens)
|
| 786 |
+
hidden_states = hidden_states + attn_out
|
| 787 |
+
|
| 788 |
+
ffn_in = self.ln2(hidden_states)
|
| 789 |
+
if keep is not None:
|
| 790 |
+
ffn_in = ffn_in * keep
|
| 791 |
+
ffn_out, ffn_shift = self.ffn(ffn_in, ffn_shift, cu_seq_lens, keep)
|
| 792 |
+
hidden_states = hidden_states + ffn_out
|
| 793 |
+
|
| 794 |
+
if state is not None:
|
| 795 |
+
# `update_recurrent_state` copies into the pre-allocated slot rather than
|
| 796 |
+
# rebinding it, so the buffers keep fixed addresses across steps -- which
|
| 797 |
+
# is what lets a captured CUDA graph replay the decode loop.
|
| 798 |
+
state.write(self.layer_id, att_shift, ffn_shift, wkv)
|
| 799 |
+
return hidden_states, v_first, state
|
| 800 |
+
|
| 801 |
+
|
| 802 |
+
@dataclass
|
| 803 |
+
class Rwkv7Output(ModelOutput):
|
| 804 |
+
r"""
|
| 805 |
+
state (`Rwkv7Cache`, *optional*):
|
| 806 |
+
The recurrent state of every block. Feed it back to continue a sequence; it
|
| 807 |
+
replaces the KV cache and is a constant size, whatever the context length.
|
| 808 |
+
attentions (always `None`):
|
| 809 |
+
Present so that code written against the common output shape does not have to
|
| 810 |
+
special-case this model. There is no attention here to report -- the mixing is
|
| 811 |
+
a recurrence, not a score matrix -- so it is `None` whatever `output_attentions`
|
| 812 |
+
is set to, rather than an empty tuple pretending to be a per-layer list.
|
| 813 |
+
"""
|
| 814 |
+
|
| 815 |
+
last_hidden_state: torch.FloatTensor | None = None
|
| 816 |
+
state: Rwkv7Cache | None = None
|
| 817 |
+
hidden_states: tuple[torch.FloatTensor, ...] | None = None
|
| 818 |
+
attentions: None = None
|
| 819 |
+
|
| 820 |
+
|
| 821 |
+
@dataclass
|
| 822 |
+
class Rwkv7CausalLMOutput(ModelOutput):
|
| 823 |
+
r"""
|
| 824 |
+
state (`Rwkv7Cache`, *optional*):
|
| 825 |
+
The recurrent state, as in [`Rwkv7Output`].
|
| 826 |
+
"""
|
| 827 |
+
|
| 828 |
+
loss: torch.FloatTensor | None = None
|
| 829 |
+
logits: torch.FloatTensor | None = None
|
| 830 |
+
state: Rwkv7Cache | None = None
|
| 831 |
+
hidden_states: tuple[torch.FloatTensor, ...] | None = None
|
| 832 |
+
attentions: None = None
|
| 833 |
+
|
| 834 |
+
|
| 835 |
+
@auto_docstring
|
| 836 |
+
class Rwkv7PreTrainedModel(PreTrainedModel):
|
| 837 |
+
config: Rwkv7Config
|
| 838 |
+
base_model_prefix = "rwkv7"
|
| 839 |
+
_no_split_modules = ["Rwkv7Block"]
|
| 840 |
+
supports_gradient_checkpointing = True
|
| 841 |
+
_is_stateful = True
|
| 842 |
+
# Beam search reorders through `Rwkv7Cache.reorder_cache`. Defining
|
| 843 |
+
# `_reorder_cache` here instead would take precedence over it in `generate` and
|
| 844 |
+
# bypass the cache's own bookkeeping, so it is deliberately absent.
|
| 845 |
+
|
| 846 |
+
@torch.no_grad()
|
| 847 |
+
def _init_weights(self, module):
|
| 848 |
+
"""Reproduce RWKV-LM's x070 constructor and final `generate_init_weight` pass."""
|
| 849 |
+
super()._init_weights(module)
|
| 850 |
+
if isinstance(module, Rwkv7Attention):
|
| 851 |
+
layer_id = module.layer_id
|
| 852 |
+
num_layers = module.config.num_hidden_layers
|
| 853 |
+
hidden_size = module.config.hidden_size
|
| 854 |
+
head_dim = module.config.head_dim
|
| 855 |
+
|
| 856 |
+
ratio_0_to_1 = layer_id / (num_layers - 1) if num_layers > 1 else 0.0
|
| 857 |
+
ratio_1_to_almost0 = 1.0 - layer_id / num_layers
|
| 858 |
+
channel = torch.arange(hidden_size, dtype=torch.float32, device=module.x_r.device)
|
| 859 |
+
ddd = (channel / hidden_size).reshape(1, 1, hidden_size)
|
| 860 |
+
linear = channel / max(hidden_size - 1, 1) - 0.5
|
| 861 |
+
if head_dim > 1:
|
| 862 |
+
zigzag = ((channel % head_dim) - (head_dim - 1) / 2) / ((head_dim - 1) / 2)
|
| 863 |
+
zigzag = zigzag * zigzag.abs()
|
| 864 |
+
else:
|
| 865 |
+
zigzag = torch.zeros_like(channel)
|
| 866 |
+
time_decay = -6 + 6 * (channel / max(hidden_size - 1, 1)) ** (1 + ratio_0_to_1**0.3)
|
| 867 |
+
|
| 868 |
+
init.copy_(module.x_r, 1 - ddd ** (0.2 * ratio_1_to_almost0))
|
| 869 |
+
init.copy_(module.x_w, 1 - ddd ** (0.9 * ratio_1_to_almost0))
|
| 870 |
+
init.copy_(module.x_k, 1 - ddd ** (0.7 * ratio_1_to_almost0))
|
| 871 |
+
init.copy_(module.x_v, 1 - ddd ** (0.7 * ratio_1_to_almost0))
|
| 872 |
+
init.copy_(module.x_a, 1 - ddd ** (0.9 * ratio_1_to_almost0))
|
| 873 |
+
init.copy_(module.x_g, 1 - ddd ** (0.2 * ratio_1_to_almost0))
|
| 874 |
+
|
| 875 |
+
init.zeros_(module.w1)
|
| 876 |
+
_rwkv7_lora_orthogonal_(module.w2)
|
| 877 |
+
init.copy_(module.w0, (time_decay + 0.5 + zigzag * 2.5).reshape(1, 1, hidden_size))
|
| 878 |
+
init.zeros_(module.a1)
|
| 879 |
+
_rwkv7_lora_orthogonal_(module.a2)
|
| 880 |
+
init.copy_(module.a0, (-0.19 + zigzag * 0.3 + linear * 0.4).reshape(1, 1, hidden_size))
|
| 881 |
+
init.zeros_(module.v1)
|
| 882 |
+
_rwkv7_lora_orthogonal_(module.v2)
|
| 883 |
+
init.copy_(module.v0, (0.73 - linear * 0.4).reshape(1, 1, hidden_size))
|
| 884 |
+
init.zeros_(module.g1)
|
| 885 |
+
_rwkv7_lora_orthogonal_(module.g2)
|
| 886 |
+
init.copy_(module.k_k, (0.71 - linear * 0.1).reshape(1, 1, hidden_size))
|
| 887 |
+
init.constant_(module.k_a, 1.02)
|
| 888 |
+
init.constant_(module.r_k, -0.04)
|
| 889 |
+
|
| 890 |
+
init.orthogonal_(module.key.weight, gain=0.1)
|
| 891 |
+
init.zeros_(module.output.weight)
|
| 892 |
+
init.constant_(module.ln_x.weight, ((layer_id + 1) / num_layers) ** 0.7)
|
| 893 |
+
if module.ln_x.bias is not None:
|
| 894 |
+
init.zeros_(module.ln_x.bias)
|
| 895 |
+
elif isinstance(module, Rwkv7FeedForward):
|
| 896 |
+
ratio_1_to_almost0 = 1.0 - module.layer_id / module.config.num_hidden_layers
|
| 897 |
+
channel = torch.arange(module.config.hidden_size, dtype=torch.float32, device=module.x_k.device)
|
| 898 |
+
ddd = (channel / module.config.hidden_size).reshape(1, 1, module.config.hidden_size)
|
| 899 |
+
init.copy_(module.x_k, 1 - ddd ** (ratio_1_to_almost0**4))
|
| 900 |
+
init.zeros_(module.value.weight)
|
| 901 |
+
elif isinstance(module, nn.Linear):
|
| 902 |
+
init.orthogonal_(module.weight)
|
| 903 |
+
if module.bias is not None:
|
| 904 |
+
init.zeros_(module.bias)
|
| 905 |
+
elif isinstance(module, nn.Embedding):
|
| 906 |
+
init.uniform_(module.weight, a=-1e-4, b=1e-4)
|
| 907 |
+
|
| 908 |
+
|
| 909 |
+
@auto_docstring
|
| 910 |
+
class Rwkv7Model(Rwkv7PreTrainedModel):
|
| 911 |
+
def __init__(self, config: Rwkv7Config):
|
| 912 |
+
super().__init__(config)
|
| 913 |
+
self.emb = nn.Embedding(config.vocab_size, config.hidden_size)
|
| 914 |
+
self.blocks = nn.ModuleList([Rwkv7Block(config, i) for i in range(config.num_hidden_layers)])
|
| 915 |
+
self.ln_out = nn.LayerNorm(config.hidden_size, eps=config.norm_eps, bias=config.norm_bias)
|
| 916 |
+
self.gradient_checkpointing = False
|
| 917 |
+
self.post_init()
|
| 918 |
+
|
| 919 |
+
def get_input_embeddings(self):
|
| 920 |
+
return self.emb
|
| 921 |
+
|
| 922 |
+
def set_input_embeddings(self, new_embeddings):
|
| 923 |
+
self.emb = new_embeddings
|
| 924 |
+
|
| 925 |
+
def allocate_state(self, batch: int, device=None, dtype=None) -> Rwkv7Cache:
|
| 926 |
+
"""A zeroed cache, plus everything else the decode path would build lazily.
|
| 927 |
+
|
| 928 |
+
Call this before compiling, and pass the result in as `state=`.
|
| 929 |
+
`mark_static_address` cannot run during tracing, so anything first allocated
|
| 930 |
+
*inside* the compiled region stays unpinned, and inductor declines CUDA graphs
|
| 931 |
+
for a region that mutates its inputs. Starting from `state=None` loses the
|
| 932 |
+
state buffers that way, and the decode then runs several times slower while
|
| 933 |
+
saying so only in a line of warning.
|
| 934 |
+
"""
|
| 935 |
+
state = self._empty_state(
|
| 936 |
+
batch,
|
| 937 |
+
device if device is not None else self.emb.weight.device,
|
| 938 |
+
dtype if dtype is not None else self.emb.weight.dtype,
|
| 939 |
+
)
|
| 940 |
+
return state
|
| 941 |
+
|
| 942 |
+
def _empty_state(self, batch: int, device, dtype) -> Rwkv7Cache:
|
| 943 |
+
return Rwkv7Cache(self.config, batch_size=batch, device=device, dtype=dtype)
|
| 944 |
+
|
| 945 |
+
@can_return_tuple
|
| 946 |
+
@auto_docstring
|
| 947 |
+
def forward(
|
| 948 |
+
self,
|
| 949 |
+
input_ids: torch.LongTensor | None = None,
|
| 950 |
+
attention_mask: torch.LongTensor | None = None,
|
| 951 |
+
inputs_embeds: torch.FloatTensor | None = None,
|
| 952 |
+
state: Rwkv7Cache | None = None,
|
| 953 |
+
cu_seq_lens: torch.LongTensor | None = None,
|
| 954 |
+
use_cache: bool | None = None,
|
| 955 |
+
output_attentions: bool | None = None,
|
| 956 |
+
output_hidden_states: bool | None = None,
|
| 957 |
+
**kwargs,
|
| 958 |
+
) -> Rwkv7Output:
|
| 959 |
+
r"""
|
| 960 |
+
attention_mask (`torch.LongTensor`, *optional*):
|
| 961 |
+
1 on real tokens, 0 on padding. Read as the **tail** of whatever is
|
| 962 |
+
given, so a decode step may hand over the whole conversation's mask and
|
| 963 |
+
only the last position is used. A prefix chunk must therefore slice its
|
| 964 |
+
own mask to match the `input_ids` it passes, or it silently masks the
|
| 965 |
+
wrong positions. There is nothing here for a mask to hide behind: an
|
| 966 |
+
all-recurrent model feeds pads through the recurrence like any other
|
| 967 |
+
token unless they are neutralised.
|
| 968 |
+
state (`Rwkv7Cache`, *optional*):
|
| 969 |
+
Recurrent state returned by a previous call; pass it back to continue
|
| 970 |
+
the sequence. Allocated on the first forward if omitted. Ignored as a
|
| 971 |
+
*history* when `cu_seq_lens` is given -- see there.
|
| 972 |
+
cu_seq_lens (`torch.LongTensor`, *optional*):
|
| 973 |
+
Cumulative sequence lengths for a *packed* batch: several sequences
|
| 974 |
+
concatenated into one row instead of padded to a rectangle, starting at
|
| 975 |
+
0 and ending at `seq_len`, and non-decreasing. Each segment then decodes
|
| 976 |
+
from a fresh recurrent state, as if it had been run on its own. This is
|
| 977 |
+
the varlen layout; use it instead of padding when the lengths vary a
|
| 978 |
+
lot, since a recurrent model pays for pad tokens in time as well as
|
| 979 |
+
memory.
|
| 980 |
+
|
| 981 |
+
A packed batch is a set of *new* sequences, not a continuation, so a
|
| 982 |
+
`state` passed alongside contributes its shape and dtype and nothing
|
| 983 |
+
else -- its contents are not read. What comes back is the last segment's
|
| 984 |
+
state, which is the one a continuation of this row would resume from.
|
| 985 |
+
"""
|
| 986 |
+
output_hidden_states = (
|
| 987 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
| 988 |
+
)
|
| 989 |
+
use_cache = use_cache if use_cache is not None else self.config.use_cache
|
| 990 |
+
# Gradient checkpointing replays the forward during backward. If the cache is
|
| 991 |
+
# live, the replay re-reads a state the first pass already advanced, and the
|
| 992 |
+
# gradients that come back are wrong rather than absent. `GradientCheckpointingLayer`
|
| 993 |
+
# neutralises this itself, but only for a cache arriving as a keyword named
|
| 994 |
+
# `use_cache` / `past_key_values` / `layer_past`; this model hands its state to
|
| 995 |
+
# the block positionally, so none of those guards see it and the whole thing
|
| 996 |
+
# has to be caught here.
|
| 997 |
+
if self.gradient_checkpointing and self.training and use_cache:
|
| 998 |
+
logger.warning_once(
|
| 999 |
+
"`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
|
| 1000 |
+
)
|
| 1001 |
+
use_cache = False
|
| 1002 |
+
|
| 1003 |
+
if (input_ids is None) == (inputs_embeds is None):
|
| 1004 |
+
raise ValueError("Specify exactly one of input_ids or inputs_embeds")
|
| 1005 |
+
if inputs_embeds is None:
|
| 1006 |
+
inputs_embeds = self.emb(input_ids)
|
| 1007 |
+
|
| 1008 |
+
if use_cache and state is None:
|
| 1009 |
+
state = self._empty_state(inputs_embeds.shape[0], inputs_embeds.device, inputs_embeds.dtype)
|
| 1010 |
+
|
| 1011 |
+
# An attention-free recurrence has to be told where the padding is: pads are
|
| 1012 |
+
# fed through the recurrence like any other token, so a left-padded batch
|
| 1013 |
+
# (what `generate` produces) would otherwise start every short row from a
|
| 1014 |
+
# state the pads had already moved. A single decoded token is never padding,
|
| 1015 |
+
# so the decode step takes none of this.
|
| 1016 |
+
keep = None
|
| 1017 |
+
if attention_mask is not None:
|
| 1018 |
+
# The mask is taken as the TAIL of whatever is passed, so a decode step
|
| 1019 |
+
# can hand over the whole conversation's mask. A prefix chunk must slice
|
| 1020 |
+
# its own mask to match, or it silently reads the wrong positions.
|
| 1021 |
+
#
|
| 1022 |
+
# Applied at seq_len == 1 as well. Skipping it there assumes a single
|
| 1023 |
+
# decoded token is never padding -- true of `generate`, but an
|
| 1024 |
+
# assumption about the caller rather than a property
|
| 1025 |
+
# of the model, and a fully-masked 1-token row was moving the state.
|
| 1026 |
+
keep = attention_mask[:, -inputs_embeds.shape[1] :, None].to(inputs_embeds.dtype)
|
| 1027 |
+
|
| 1028 |
+
if cu_seq_lens is not None:
|
| 1029 |
+
# Checked rather than trusted: a malformed boundary list does not fail,
|
| 1030 |
+
# it silently splits the recurrence in the wrong places and returns
|
| 1031 |
+
# fluent output computed from the wrong states.
|
| 1032 |
+
if inputs_embeds.shape[0] != 1:
|
| 1033 |
+
raise ValueError(
|
| 1034 |
+
f"cu_seq_lens describes one packed row, but got batch size {inputs_embeds.shape[0]}. "
|
| 1035 |
+
"Pack the sequences into a single row, or use attention_mask with a padded batch."
|
| 1036 |
+
)
|
| 1037 |
+
if cu_seq_lens.ndim != 1 or cu_seq_lens[0] != 0 or cu_seq_lens[-1] != inputs_embeds.shape[1]:
|
| 1038 |
+
raise ValueError(
|
| 1039 |
+
f"cu_seq_lens must be 1-D, start at 0 and end at seq_len ({inputs_embeds.shape[1]}); "
|
| 1040 |
+
f"got {cu_seq_lens.tolist()}"
|
| 1041 |
+
)
|
| 1042 |
+
# Endpoints alone do not pin the list down. A pair that goes backwards
|
| 1043 |
+
# is skipped rather than rejected further in, so the segments emit fewer
|
| 1044 |
+
# tokens than came in and the row silently changes length -- and one that
|
| 1045 |
+
# merely repeats a boundary contributes an empty segment, which is
|
| 1046 |
+
# harmless but is never what the caller meant.
|
| 1047 |
+
if bool((cu_seq_lens[1:] <= cu_seq_lens[:-1]).any()):
|
| 1048 |
+
raise ValueError(
|
| 1049 |
+
"cu_seq_lens must be strictly increasing (each segment needs at least one token); "
|
| 1050 |
+
f"got {cu_seq_lens.tolist()}"
|
| 1051 |
+
)
|
| 1052 |
+
|
| 1053 |
+
hidden_states = inputs_embeds
|
| 1054 |
+
v_first = None
|
| 1055 |
+
all_hidden_states = () if output_hidden_states else None
|
| 1056 |
+
|
| 1057 |
+
for block in self.blocks:
|
| 1058 |
+
if output_hidden_states:
|
| 1059 |
+
all_hidden_states = all_hidden_states + (hidden_states,)
|
| 1060 |
+
hidden_states, v_first, state = block(hidden_states, v_first, state, keep, cu_seq_lens)
|
| 1061 |
+
|
| 1062 |
+
hidden_states = self.ln_out(hidden_states)
|
| 1063 |
+
if output_hidden_states:
|
| 1064 |
+
all_hidden_states = all_hidden_states + (hidden_states,)
|
| 1065 |
+
|
| 1066 |
+
return Rwkv7Output(last_hidden_state=hidden_states, state=state, hidden_states=all_hidden_states)
|
| 1067 |
+
|
| 1068 |
+
|
| 1069 |
+
@auto_docstring
|
| 1070 |
+
class Rwkv7ForCausalLM(Rwkv7PreTrainedModel, GenerationMixin):
|
| 1071 |
+
_tied_weights_keys = {"head.weight": "rwkv7.emb.weight"}
|
| 1072 |
+
|
| 1073 |
+
def __init__(self, config: Rwkv7Config):
|
| 1074 |
+
super().__init__(config)
|
| 1075 |
+
self.rwkv7 = Rwkv7Model(config)
|
| 1076 |
+
self.head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
| 1077 |
+
self.post_init()
|
| 1078 |
+
|
| 1079 |
+
@torch.no_grad()
|
| 1080 |
+
def _init_weights(self, module):
|
| 1081 |
+
super()._init_weights(module)
|
| 1082 |
+
# Composite-model dispatch sends the base RWKV7 modules through their own
|
| 1083 |
+
# initializer, so every Linear reaching this override is an LM head. Using
|
| 1084 |
+
# its actual shape also covers a newly resized, not-yet-installed head.
|
| 1085 |
+
if isinstance(module, nn.Linear):
|
| 1086 |
+
gain = 0.5 * math.sqrt(module.out_features / module.in_features)
|
| 1087 |
+
if module.out_features <= module.in_features:
|
| 1088 |
+
gain = 0.5
|
| 1089 |
+
init.orthogonal_(module.weight, gain=gain)
|
| 1090 |
+
|
| 1091 |
+
def get_output_embeddings(self):
|
| 1092 |
+
return self.head
|
| 1093 |
+
|
| 1094 |
+
def set_output_embeddings(self, new_embeddings):
|
| 1095 |
+
self.head = new_embeddings
|
| 1096 |
+
|
| 1097 |
+
def prepare_inputs_for_generation(
|
| 1098 |
+
self,
|
| 1099 |
+
input_ids,
|
| 1100 |
+
state=None,
|
| 1101 |
+
inputs_embeds=None,
|
| 1102 |
+
next_sequence_length=None,
|
| 1103 |
+
is_first_iteration=False,
|
| 1104 |
+
**kwargs,
|
| 1105 |
+
):
|
| 1106 |
+
# `state is not None` does not by itself mean decode: callers can provide
|
| 1107 |
+
# an empty preallocated state for the initial prompt, or a carried state
|
| 1108 |
+
# followed by a multi-token continuation. GenerationMixin tells us how
|
| 1109 |
+
# much of the input is new through `next_sequence_length`.
|
| 1110 |
+
if next_sequence_length is not None:
|
| 1111 |
+
input_ids = input_ids[:, -next_sequence_length:]
|
| 1112 |
+
if inputs_embeds is not None:
|
| 1113 |
+
inputs_embeds = inputs_embeds[:, -next_sequence_length:]
|
| 1114 |
+
elif state is not None and not is_first_iteration:
|
| 1115 |
+
input_ids = input_ids[:, -1:]
|
| 1116 |
+
model_inputs = (
|
| 1117 |
+
{"input_ids": input_ids}
|
| 1118 |
+
if inputs_embeds is None or not is_first_iteration
|
| 1119 |
+
else {"inputs_embeds": inputs_embeds}
|
| 1120 |
+
)
|
| 1121 |
+
model_inputs["state"] = state
|
| 1122 |
+
# Everything else the caller passed goes through, minus what this model does
|
| 1123 |
+
# not take. An allowlist of specific names instead would make
|
| 1124 |
+
# `generate(output_hidden_states=True)` return a tuple of `None`: the flag
|
| 1125 |
+
# would be dropped here, the forward would never see it, and generate collects the
|
| 1126 |
+
# nothing it got back. Any user kwarg met the same fate, silently. The two
|
| 1127 |
+
# excluded here are `labels`, which would make generate compute a loss it
|
| 1128 |
+
# never reads, and the KV-cache bookkeeping that belongs to models with a KV
|
| 1129 |
+
# cache; this one carries its history in `state`.
|
| 1130 |
+
skip = ("labels", "past_key_values", "cache_position")
|
| 1131 |
+
model_inputs.update({k: v for k, v in kwargs.items() if k not in skip and k not in model_inputs})
|
| 1132 |
+
return model_inputs
|
| 1133 |
+
|
| 1134 |
+
@can_return_tuple
|
| 1135 |
+
@auto_docstring
|
| 1136 |
+
def forward(
|
| 1137 |
+
self,
|
| 1138 |
+
input_ids: torch.LongTensor | None = None,
|
| 1139 |
+
attention_mask: torch.LongTensor | None = None,
|
| 1140 |
+
inputs_embeds: torch.FloatTensor | None = None,
|
| 1141 |
+
state: Rwkv7Cache | None = None,
|
| 1142 |
+
labels: torch.LongTensor | None = None,
|
| 1143 |
+
use_cache: bool | None = None,
|
| 1144 |
+
output_attentions: bool | None = None,
|
| 1145 |
+
output_hidden_states: bool | None = None,
|
| 1146 |
+
logits_to_keep: int | torch.Tensor = 0,
|
| 1147 |
+
**kwargs,
|
| 1148 |
+
) -> Rwkv7CausalLMOutput:
|
| 1149 |
+
r"""
|
| 1150 |
+
state (`Rwkv7Cache`, *optional*):
|
| 1151 |
+
Recurrent state returned by a previous call.
|
| 1152 |
+
logits_to_keep (`int` or `torch.Tensor`, *optional*, defaults to 0):
|
| 1153 |
+
Compute the head over only the last `logits_to_keep` positions, or over the
|
| 1154 |
+
positions this tensor indexes; 0 means all of them. Worth having on a
|
| 1155 |
+
recurrent model for the same reason as on any other: a prefill needs one
|
| 1156 |
+
row of logits and the vocabulary is the widest matrix in the model, so
|
| 1157 |
+
running the head over the whole prompt is the single largest avoidable cost
|
| 1158 |
+
in a prefill. Until this argument existed it was swallowed by `**kwargs`,
|
| 1159 |
+
so `generate` declined to pass it and a caller who passed it was quietly
|
| 1160 |
+
ignored.
|
| 1161 |
+
"""
|
| 1162 |
+
outputs = self.rwkv7(
|
| 1163 |
+
input_ids=input_ids,
|
| 1164 |
+
attention_mask=attention_mask,
|
| 1165 |
+
inputs_embeds=inputs_embeds,
|
| 1166 |
+
state=state,
|
| 1167 |
+
use_cache=use_cache,
|
| 1168 |
+
output_hidden_states=output_hidden_states,
|
| 1169 |
+
**kwargs,
|
| 1170 |
+
)
|
| 1171 |
+
keep = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
|
| 1172 |
+
hidden_to_project = outputs.last_hidden_state[:, keep, :]
|
| 1173 |
+
logits = _rwkv7_linear(
|
| 1174 |
+
self.head,
|
| 1175 |
+
hidden_to_project,
|
| 1176 |
+
self.config.wkv_implementation != "chunked",
|
| 1177 |
+
)
|
| 1178 |
+
|
| 1179 |
+
loss = None
|
| 1180 |
+
if labels is not None:
|
| 1181 |
+
loss = self.loss_function(logits, labels, self.config.vocab_size, **kwargs)
|
| 1182 |
+
|
| 1183 |
+
return Rwkv7CausalLMOutput(loss=loss, logits=logits, state=outputs.state, hidden_states=outputs.hidden_states)
|
| 1184 |
+
|
| 1185 |
+
|
| 1186 |
+
__all__ = ["Rwkv7Cache", "Rwkv7PreTrainedModel", "Rwkv7Model", "Rwkv7ForCausalLM"]
|
release-manifest.json
ADDED
|
@@ -0,0 +1,332 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"builder": {
|
| 3 |
+
"asset_set": "2026.08.06-r1",
|
| 4 |
+
"assets_sha256": "3c3f6237fcd9dc2ec8ba0b61b626ef72924ebe42bdc6dc49e21a5bfeb19bb96b",
|
| 5 |
+
"version": "0.2.0"
|
| 6 |
+
},
|
| 7 |
+
"conversion": {
|
| 8 |
+
"explicit_cast": false,
|
| 9 |
+
"serialized_parameter_count": 1527668736,
|
| 10 |
+
"source_float_dtypes": [
|
| 11 |
+
"bfloat16"
|
| 12 |
+
],
|
| 13 |
+
"source_parameter_count": 1527668736,
|
| 14 |
+
"synthesized_tensors": [],
|
| 15 |
+
"target_float_dtypes": [
|
| 16 |
+
"bfloat16"
|
| 17 |
+
],
|
| 18 |
+
"tensor_count": 798,
|
| 19 |
+
"tensor_map_sha256": "03131ced241e7b0f86869363b3969ead6ef58462f23d9f04e56e26d00565aefc"
|
| 20 |
+
},
|
| 21 |
+
"files": {
|
| 22 |
+
".gitattributes": {
|
| 23 |
+
"role": "metadata",
|
| 24 |
+
"sha256": "edfb97afcbabd3c89eb1d19070521d403f4644ecacfec826e14c2dc68cf419d8",
|
| 25 |
+
"size_bytes": 134
|
| 26 |
+
},
|
| 27 |
+
"LICENSE": {
|
| 28 |
+
"role": "metadata",
|
| 29 |
+
"sha256": "475641f33cfd755f19245d2eb6a6d8fb7fdca60835cfa369b7e3774460409cc4",
|
| 30 |
+
"size_bytes": 11357
|
| 31 |
+
},
|
| 32 |
+
"NOTICE": {
|
| 33 |
+
"role": "metadata",
|
| 34 |
+
"sha256": "7f8e7f0f3bf0a8e7caf13132770cfe841e5eec10203c18798f8051e96d234108",
|
| 35 |
+
"size_bytes": 250
|
| 36 |
+
},
|
| 37 |
+
"README.md": {
|
| 38 |
+
"role": "model_card",
|
| 39 |
+
"sha256": "b7795213e1fcf1d9e7de41fa027a0b2d311b2dacfa5a5ba9f62c4c2484bd729f",
|
| 40 |
+
"size_bytes": 10260
|
| 41 |
+
},
|
| 42 |
+
"chat_template.jinja": {
|
| 43 |
+
"role": "tokenizer",
|
| 44 |
+
"sha256": "f0d8f8161f03d38d0bf68b1ea9a1293064c0fe2b60dfc2874008154c33861c07",
|
| 45 |
+
"size_bytes": 2333
|
| 46 |
+
},
|
| 47 |
+
"config.json": {
|
| 48 |
+
"role": "model_config",
|
| 49 |
+
"sha256": "442d1155144fd8331bac86250131cc1a874d8d41c1504f930e0b522a99b615aa",
|
| 50 |
+
"size_bytes": 748
|
| 51 |
+
},
|
| 52 |
+
"configuration_rwkv7.py": {
|
| 53 |
+
"role": "model_code",
|
| 54 |
+
"sha256": "a4a2adcd02a101cc2a17f2452f8e56f4a8d8348be0d19ab989e908a83921c123",
|
| 55 |
+
"size_bytes": 7318
|
| 56 |
+
},
|
| 57 |
+
"generation_config.json": {
|
| 58 |
+
"role": "model_config",
|
| 59 |
+
"sha256": "358114ae91c5eee6fe10b41a22a9ee95a9b1c20ed86d0666617aa049386d63a3",
|
| 60 |
+
"size_bytes": 87
|
| 61 |
+
},
|
| 62 |
+
"inference/generate.py": {
|
| 63 |
+
"role": "inference",
|
| 64 |
+
"sha256": "9e4f3461ade971a96a276ec4e18683d6d8df6420b09bf3e63cc499b6e119aa05",
|
| 65 |
+
"size_bytes": 7804
|
| 66 |
+
},
|
| 67 |
+
"inference/kernel.py": {
|
| 68 |
+
"role": "inference",
|
| 69 |
+
"sha256": "88379d8979b908a01af1f06b89d8fc624d332796db6044ec4a027d7cecab1cb6",
|
| 70 |
+
"size_bytes": 165253
|
| 71 |
+
},
|
| 72 |
+
"inference/model_loader.py": {
|
| 73 |
+
"role": "inference",
|
| 74 |
+
"sha256": "c48abe28abde91ede570498372cd6b9aa6a05080525df7f5045613849ed9bb61",
|
| 75 |
+
"size_bytes": 6452
|
| 76 |
+
},
|
| 77 |
+
"inference/requirements.txt": {
|
| 78 |
+
"role": "inference",
|
| 79 |
+
"sha256": "9cc4f89e2edc45da7a04cdf8e4c1da83b682a6bf969e0eed7b3dec91128d925a",
|
| 80 |
+
"size_bytes": 93
|
| 81 |
+
},
|
| 82 |
+
"inference/runtime.py": {
|
| 83 |
+
"role": "inference",
|
| 84 |
+
"sha256": "064c53ce5dfe811ea873ad9976f8f0839b55e515e604adf3c3e28fd4c02de784",
|
| 85 |
+
"size_bytes": 158191
|
| 86 |
+
},
|
| 87 |
+
"model.safetensors": {
|
| 88 |
+
"role": "weights",
|
| 89 |
+
"sha256": "84ccbb857c84e00cefc48b233937ada79c411e491df25fb21aed23237f39a14f",
|
| 90 |
+
"size_bytes": 3055418240
|
| 91 |
+
},
|
| 92 |
+
"modeling_rwkv7.py": {
|
| 93 |
+
"role": "model_code",
|
| 94 |
+
"sha256": "f3a1d9bde755f54887451e8755c6d48866ecc9146610c7d8fa73e0fa68c6734c",
|
| 95 |
+
"size_bytes": 57446
|
| 96 |
+
},
|
| 97 |
+
"tokenizer.json": {
|
| 98 |
+
"role": "tokenizer",
|
| 99 |
+
"sha256": "1c0bc8e5f355de60a4587a1e9abb9aa4ef65508ff59641344e7f9f10f76feacf",
|
| 100 |
+
"size_bytes": 1580973
|
| 101 |
+
},
|
| 102 |
+
"tokenizer_config.json": {
|
| 103 |
+
"role": "tokenizer",
|
| 104 |
+
"sha256": "58a994975a1a64327cd7ee71f230d4a5f290e7260957b18840e9ecf0d5d57d0e",
|
| 105 |
+
"size_bytes": 407
|
| 106 |
+
}
|
| 107 |
+
},
|
| 108 |
+
"identity": {
|
| 109 |
+
"parameter_label": "1.5",
|
| 110 |
+
"release_date": "20260805",
|
| 111 |
+
"training_context_length": 16384
|
| 112 |
+
},
|
| 113 |
+
"metadata": {
|
| 114 |
+
"context_length": 16384,
|
| 115 |
+
"datasets": [
|
| 116 |
+
"HuggingFaceFW/fineweb-edu",
|
| 117 |
+
"mlfoundations/dclm-baseline-1.0",
|
| 118 |
+
"cerebras/SlimPajama-627B",
|
| 119 |
+
"EleutherAI/pile",
|
| 120 |
+
"bigcode/starcoderdata",
|
| 121 |
+
"oscar-corpus/OSCAR-2301"
|
| 122 |
+
],
|
| 123 |
+
"languages": [
|
| 124 |
+
"en",
|
| 125 |
+
"zh",
|
| 126 |
+
"fr",
|
| 127 |
+
"es",
|
| 128 |
+
"de",
|
| 129 |
+
"pt",
|
| 130 |
+
"ru",
|
| 131 |
+
"it",
|
| 132 |
+
"ja",
|
| 133 |
+
"ko",
|
| 134 |
+
"vi",
|
| 135 |
+
"ar"
|
| 136 |
+
],
|
| 137 |
+
"license": "apache-2.0",
|
| 138 |
+
"profile": "g1i",
|
| 139 |
+
"provenance": "locked-profile"
|
| 140 |
+
},
|
| 141 |
+
"model_code": {
|
| 142 |
+
"format_version": 1,
|
| 143 |
+
"source_repository": "https://github.com/huggingface/transformers.git",
|
| 144 |
+
"source_revision": "4ad9ed0747ed6ba75c787e8f9040dcd64b166ee2",
|
| 145 |
+
"sources": {
|
| 146 |
+
"configuration_rwkv7.py": {
|
| 147 |
+
"asset_path": "model_code/configuration_rwkv7.py",
|
| 148 |
+
"output_sha256": "a4a2adcd02a101cc2a17f2452f8e56f4a8d8348be0d19ab989e908a83921c123",
|
| 149 |
+
"repository_path": "src/transformers/models/rwkv7/configuration_rwkv7.py",
|
| 150 |
+
"source_sha256": "6f5b92c5fe7498ad22b0054a2f735a7ca82e7577436f4ad32f0fc27d1e900fdd"
|
| 151 |
+
},
|
| 152 |
+
"modeling_rwkv7.py": {
|
| 153 |
+
"asset_path": "model_code/modeling_rwkv7.py",
|
| 154 |
+
"output_sha256": "f3a1d9bde755f54887451e8755c6d48866ecc9146610c7d8fa73e0fa68c6734c",
|
| 155 |
+
"repository_path": "src/transformers/models/rwkv7/modeling_rwkv7.py",
|
| 156 |
+
"source_sha256": "3e8e5af7c4eba0b5de1496aef44773d7ac1bb4d96756e6f55efaf29453d67952"
|
| 157 |
+
}
|
| 158 |
+
},
|
| 159 |
+
"transformers_min_version": "5.15"
|
| 160 |
+
},
|
| 161 |
+
"profile": {
|
| 162 |
+
"checkpoint": "g1i-1.5b-20260805",
|
| 163 |
+
"family": "g1i"
|
| 164 |
+
},
|
| 165 |
+
"runtime": {
|
| 166 |
+
"files": {
|
| 167 |
+
"kernel.py": {
|
| 168 |
+
"sha256": "88379d8979b908a01af1f06b89d8fc624d332796db6044ec4a027d7cecab1cb6"
|
| 169 |
+
},
|
| 170 |
+
"runtime.py": {
|
| 171 |
+
"sha256": "064c53ce5dfe811ea873ad9976f8f0839b55e515e604adf3c3e28fd4c02de784"
|
| 172 |
+
}
|
| 173 |
+
},
|
| 174 |
+
"format_version": 4,
|
| 175 |
+
"kernel": {
|
| 176 |
+
"format_version": 3,
|
| 177 |
+
"output": "inference/kernel.py",
|
| 178 |
+
"output_sha256": "88379d8979b908a01af1f06b89d8fc624d332796db6044ec4a027d7cecab1cb6",
|
| 179 |
+
"sources": {
|
| 180 |
+
"kernel_tilelang_decode.py": {
|
| 181 |
+
"exports": [
|
| 182 |
+
"Any",
|
| 183 |
+
"_HEAD_SIZE",
|
| 184 |
+
"_NEXP_HALF_LOG2_E",
|
| 185 |
+
"_NLOG2_E",
|
| 186 |
+
"_ROTATOR1_SIGNED",
|
| 187 |
+
"_TWO_NEG_41",
|
| 188 |
+
"_build_cmix_add_layernorm_mix_program",
|
| 189 |
+
"_build_cmix_binned_finalize_program",
|
| 190 |
+
"_build_cmix_finalize_program",
|
| 191 |
+
"_build_cmix_layernorm_mix_program",
|
| 192 |
+
"_build_cmix_sparse_atomic_program",
|
| 193 |
+
"_build_cmix_sparse_binned_program",
|
| 194 |
+
"_build_cmix_sparse_split_program",
|
| 195 |
+
"_build_cmix_value_program",
|
| 196 |
+
"_build_ffn_program",
|
| 197 |
+
"_build_gemv_program",
|
| 198 |
+
"_build_key_gate_program",
|
| 199 |
+
"_build_post_state_program",
|
| 200 |
+
"_build_rankout_program",
|
| 201 |
+
"_build_rankout_reduced_program",
|
| 202 |
+
"_build_rkv_program",
|
| 203 |
+
"_build_tmix_layernorm_mix6_program",
|
| 204 |
+
"_build_wkv_program",
|
| 205 |
+
"_build_wkv_w0_t1_program",
|
| 206 |
+
"_cmix_value_out",
|
| 207 |
+
"_compiled_cmix_add_layernorm_mix",
|
| 208 |
+
"_compiled_cmix_binned_finalize",
|
| 209 |
+
"_compiled_cmix_finalize",
|
| 210 |
+
"_compiled_cmix_layernorm_mix",
|
| 211 |
+
"_compiled_cmix_sparse_atomic",
|
| 212 |
+
"_compiled_cmix_sparse_binned",
|
| 213 |
+
"_compiled_cmix_sparse_split",
|
| 214 |
+
"_compiled_cmix_value",
|
| 215 |
+
"_compiled_ffn",
|
| 216 |
+
"_compiled_gemv",
|
| 217 |
+
"_compiled_key_gate",
|
| 218 |
+
"_compiled_post_state",
|
| 219 |
+
"_compiled_rankout",
|
| 220 |
+
"_compiled_rankout_reduced",
|
| 221 |
+
"_compiled_rkv",
|
| 222 |
+
"_compiled_tmix_layernorm_mix6",
|
| 223 |
+
"_compiled_wkv",
|
| 224 |
+
"_compiled_wkv_w0_t1",
|
| 225 |
+
"_dtype_name",
|
| 226 |
+
"_ffn_out",
|
| 227 |
+
"_post_state_out",
|
| 228 |
+
"_rankout_out",
|
| 229 |
+
"_require_contiguous_cuda",
|
| 230 |
+
"_rkv_out",
|
| 231 |
+
"_wkv_kernel_out",
|
| 232 |
+
"_wkv_out",
|
| 233 |
+
"_wkv_precise_out",
|
| 234 |
+
"_wkv_w0_t1_out",
|
| 235 |
+
"clear_tilelang_kernel_caches",
|
| 236 |
+
"cuda_arch_key",
|
| 237 |
+
"lru_cache",
|
| 238 |
+
"torch"
|
| 239 |
+
],
|
| 240 |
+
"namespace": "decode",
|
| 241 |
+
"sha256": "b582151a6d2bd345cc5f87415543370e027c8f943a4dfe57a18300f08560bd73"
|
| 242 |
+
},
|
| 243 |
+
"kernel_tilelang_state.py": {
|
| 244 |
+
"exports": [
|
| 245 |
+
"Any",
|
| 246 |
+
"BFLOAT16",
|
| 247 |
+
"CUDA_GET_DEVICE_CAPABILITY",
|
| 248 |
+
"EXACT_FUSED_STATE_PROJECTION_CAPABILITIES",
|
| 249 |
+
"FLOAT16",
|
| 250 |
+
"FLOAT32",
|
| 251 |
+
"IS_GRAD_ENABLED",
|
| 252 |
+
"TORCH_STACK",
|
| 253 |
+
"TORCH_WHERE",
|
| 254 |
+
"TORCH_ZEROS_LIKE",
|
| 255 |
+
"_compiled_backward_kernel",
|
| 256 |
+
"_compiled_fast_state_scan",
|
| 257 |
+
"_compiled_fused_state_update",
|
| 258 |
+
"_compiled_kernel",
|
| 259 |
+
"_compiled_low_precision_state_kernel",
|
| 260 |
+
"_compiled_post_state",
|
| 261 |
+
"_compiled_state_projection",
|
| 262 |
+
"_compiled_x_mix",
|
| 263 |
+
"_dtype_name",
|
| 264 |
+
"_require_exact_fused_state_projection",
|
| 265 |
+
"build_fast_state_scan_program",
|
| 266 |
+
"build_fused_state_update_program",
|
| 267 |
+
"build_low_precision_state_program",
|
| 268 |
+
"build_post_state_program",
|
| 269 |
+
"build_state_backward_program",
|
| 270 |
+
"build_state_program",
|
| 271 |
+
"build_state_projection_program",
|
| 272 |
+
"build_x_mix_program",
|
| 273 |
+
"clear_tilelang_state_kernel_caches",
|
| 274 |
+
"cuda_arch_key",
|
| 275 |
+
"exact_fused_state_projection_supported",
|
| 276 |
+
"lru_cache",
|
| 277 |
+
"tilelang_fast_state_scan",
|
| 278 |
+
"tilelang_fused_state_update",
|
| 279 |
+
"tilelang_low_precision_state_finalize",
|
| 280 |
+
"tilelang_post_state",
|
| 281 |
+
"tilelang_state_finalize",
|
| 282 |
+
"tilelang_state_finalize_backward",
|
| 283 |
+
"tilelang_state_projection",
|
| 284 |
+
"tilelang_state_scan",
|
| 285 |
+
"tilelang_state_update",
|
| 286 |
+
"tilelang_x_mix",
|
| 287 |
+
"torch",
|
| 288 |
+
"torch_state_scan",
|
| 289 |
+
"torch_state_update"
|
| 290 |
+
],
|
| 291 |
+
"namespace": "state",
|
| 292 |
+
"sha256": "a3597160cb8d7317151200d05b241af7edcad2928c3731687764d914c28bf21a"
|
| 293 |
+
}
|
| 294 |
+
}
|
| 295 |
+
},
|
| 296 |
+
"transforms": {
|
| 297 |
+
"configuration_rwkv7.py": {
|
| 298 |
+
"output_sha256": "60a060199526f5a6c4d5a2436520eed97e9f21f2ea3cb6e6584d9b5fcb6f21f4",
|
| 299 |
+
"source_sha256": "60a060199526f5a6c4d5a2436520eed97e9f21f2ea3cb6e6584d9b5fcb6f21f4"
|
| 300 |
+
},
|
| 301 |
+
"custom_ops.py": {
|
| 302 |
+
"output_sha256": "8240d4eebc9decbd5ebf8fbcc3abc0ee6441b0eb4854b8d7bc18172cd2babf17",
|
| 303 |
+
"source_sha256": "fb86e5c9dfc9deb5a1be821ce02c3a390ce5c95ff55feded6add2422f0dc3623"
|
| 304 |
+
},
|
| 305 |
+
"kernel_dispatch.py": {
|
| 306 |
+
"output_sha256": "0290559437016ba4e2c2a26f7c7a0985b6fbab991535b7ae65ba18eb421fa9e1",
|
| 307 |
+
"source_sha256": "0290559437016ba4e2c2a26f7c7a0985b6fbab991535b7ae65ba18eb421fa9e1"
|
| 308 |
+
},
|
| 309 |
+
"modeling_rwkv7.py": {
|
| 310 |
+
"output_sha256": "ad7f8c981d21a3b0361cd26a2c68234a1bfbacc31fe52b49c3fe9d79e5b32844",
|
| 311 |
+
"source_sha256": "7e1e57a4516efe78a74f971ada0fe58cb71c87102479ed7516b6d3585964ca21"
|
| 312 |
+
},
|
| 313 |
+
"state.py": {
|
| 314 |
+
"output_sha256": "d13d81af6a301a70024494f03b95006bf5f2471ba0b57fd7f1f94bd54f70f2ef",
|
| 315 |
+
"source_sha256": "d13d81af6a301a70024494f03b95006bf5f2471ba0b57fd7f1f94bd54f70f2ef"
|
| 316 |
+
},
|
| 317 |
+
"tilelang_decode.py": {
|
| 318 |
+
"output_sha256": "313120b38178642c17ef16531c4510de87dc0f0d79d80f70b6ade8818b4dfbe7",
|
| 319 |
+
"source_sha256": "fe63e6332037cacbd0400d82f65fd6f1c282e31d92906c5f34d39eb9e83eb642"
|
| 320 |
+
}
|
| 321 |
+
}
|
| 322 |
+
},
|
| 323 |
+
"schema_version": 6,
|
| 324 |
+
"source": {
|
| 325 |
+
"filename": "rwkv7-g1i-1.5b-20260805-ctx16384.pth",
|
| 326 |
+
"kind": "huggingface",
|
| 327 |
+
"reference": "BlinkDL/rwkv7-g1/rwkv7-g1i-1.5b-20260805-ctx16384.pth",
|
| 328 |
+
"revision": "ede85bf8ab2e59aff7d7ca909fbbc73317866d89",
|
| 329 |
+
"sha256": "32ef7b5bf4dc8bde843cf26dfad809a1f527e2e76a9e790e7d406e71bcd785da",
|
| 330 |
+
"size_bytes": 3055444605
|
| 331 |
+
}
|
| 332 |
+
}
|
tokenizer.json
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
tokenizer_config.json
ADDED
|
@@ -0,0 +1,18 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"added_tokens_decoder": {
|
| 3 |
+
"0": {
|
| 4 |
+
"content": "<|endoftext|>",
|
| 5 |
+
"lstrip": false,
|
| 6 |
+
"normalized": false,
|
| 7 |
+
"rstrip": false,
|
| 8 |
+
"single_word": false,
|
| 9 |
+
"special": true
|
| 10 |
+
}
|
| 11 |
+
},
|
| 12 |
+
"backend": "tokenizers",
|
| 13 |
+
"eos_token": "<|endoftext|>",
|
| 14 |
+
"pad_token": "<|endoftext|>",
|
| 15 |
+
"padding_side": "left",
|
| 16 |
+
"tokenizer_class": "PreTrainedTokenizerFast",
|
| 17 |
+
"unk_token": "<|endoftext|>"
|
| 18 |
+
}
|