rwkv / rwkv-onnx Goto Github PK

View Code? Open in Web Editor NEW

40.0 2.0 3.0 7.15 MB

A converter and basic tester for rwkv onnx

Python 99.76% Shell 0.24%

rwkv-onnx's Introduction

rwkv-onnx

A converter and basic tester for rwkv V5 onnx

supports

fp16, fp32
onnx opsversions 15/17/18
file sizes > 2GB
cuda, cpu, tensorRT, mps

Note: converter requires protobuf 3.20.0

Install with pip install protobuf==3.20.0

rwkv-onnx's People

Contributors

Stargazers

Watchers

Forkers

vitaly-z stjordanis frank05080

rwkv-onnx's Issues

[BUG] Segmentation Fault when convert RWKV-3B

I locate the bug:

https://github.com/harrisonvanderbyl/rwkv-onnx/blob/5a601c7b254d20d7d9f4019083ff861f2d3df0af/opslist.py#L379-L383

When I converted the 3B model, the size of graph_def was larger than 3GB. Onnx couldn't deal with files larger than 2GB. Could we output RWKV layer by layer?

how to convert rwkv-5 model?

i used the convert.py to convert the rwkv-5 model ,when i used the test.py to test if it was successful,i failed.

[?] Select model: 
 > RWKV_24_2048_32_15.onnx

['AzureExecutionProvider', 'CPUExecutionProvider']
[?] Select execution providers(use space bar to select checkboxes): 
 > [ ] AzureExecutionProvider
   [ ] CPUExecutionProvider

[]
2024-05-18 15:30:52.673417593 [W:onnxruntime:, graph.cc:3593 CleanUnusedInitializersAndNodeArgs] Removing initializer 'PreTrainedTensor_4'. It is not used by any node and should be removed from the model.
2024-05-18 15:30:52.673432693 [W:onnxruntime:, graph.cc:3593 CleanUnusedInitializersAndNodeArgs] Removing initializer 'PreTrainedTensor_1'. It is not used by any node and should be removed from the model.
2024-05-18 15:30:52.673435491 [W:onnxruntime:, graph.cc:3593 CleanUnusedInitializersAndNodeArgs] Removing initializer 'PreTrainedTensor_0'. It is not used by any node and should be removed from the model.
tensor(float) 2048 24
(48, 2048)
(24, 40, 64, 64)
/.local/lib/python3.10/site-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.
  warnings.warn(
Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.
  0%|                                                    | 0/19 [00:00<?, ?it/s]
Traceback (most recent call last):
  File "rwkv-onnx/ds/../test.py", line 120, in <module>
    logits, state, state2 = model.forward(token,state, state2)
  File "rwkv-onnx/ds/../test.py", line 66, in forward
    outputs = sess.run(output_names, inputs)
  File "/.conda/envs/py10/lib/python3.10/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 220, in run
    return self._sess.run(output_names, input_feed, run_options)
onnxruntime.capi.onnxruntime_pybind11_state.InvalidArgument: [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Got invalid dimensions for input: instatewkv9 for the following indices
 index: 0 Got: 40 Expected: 32
 Please fix either the inputs/outputs or the model.

The convert.py may not fit the rwkv-4-world-169m.

I have tried to convert the weight file which i download from https://huggingface.co/RWKV/rwkv-4-world-169m with convert.py. I found the following problems.

All of the keys of weight file begin with 'rwkv.' compared to convert.py.
The .att. in convert.py should be replaced with .attention..
The value of rwkv.blocks.0.attention.time_decay is a 1d tensor. The headsnume and headsize in convert.py can't unpacke from w[f"rwkv.blocks.0.attention.time_decay"].shape

I don't know how to fix the third problem.

Format input state/output state

The input state/out state is separated, I add some ops in rwkv-cpp-server to concat and reshape the state tensors so that every rwkv onnx model has two input tensors and two output tensors.