RWKV-4 169m/430m in browser with ORT Web / TF.js / tfjs-tflite?

[josephrocca]

Hi, really exciting project! I'm wondering if you've published the model conversion script that you used to create the js_models files from the .pth model file? It would be awesome to see how the larger and newer models like RWKV-4 169m/430m perform in the browser! I think the inference speed of RWKV opens up many new possibilities for language models on the web.

[AXKuhta]

Exporting to ONNX is something that I've been tinkering with and I can report that the 169m RWKV-4 model does run in browser. Here's my code: https://github.com/AXKuhta/RWKV-LM/tree/onnx

There's two things missing:

• JavaScript implementation of the tokenizer
• JavaScript implementation of sample_logits().

Running python -i -u export_onnx.py and then rnn_export() will export the model as rwkw.onnx, which can then be tested with test_onnx.py and loaded from index.html. The demo in index.html uses greedy sampling and you just sorta have to visit https://goose.ai/tokenizer in order to encode/decode the text. It works using the wasm backend, but unfortunately throws an error if you try the webgl backend.

[BlinkDL]

Exporting to ONNX is something that I've been tinkering with and I can report that the 169m RWKV-4 model does run in browser. Here's my code: https://github.com/AXKuhta/RWKV-LM/tree/onnx

Great work :)

Did you get this error with webgl? cannot resolve operator 'Max' with opsets: ai.onnx v13

You can remove RWKV_HEAD_QK and RWKV-ffnPre which are not required for Pile models, and probably that will fix it.

p.s. upgrade onnxruntime to latest version and then you can test CUDAExecutionProvider in python. I think you might be using an older onnxruntime because all new versions require explicitly setting providers when initializing InferenceSession().

[josephrocca]

@AXKuhta Nice! I got a web demo going here (for 169m and 430m):

• Demo: https://josephrocca.github.io/rwkv-v4-web/demo/
• Code: https://github.com/josephrocca/rwkv-v4-web

But it seems like something is going wrong - the model isn't "coherent" in using the context. For example, if you prompt the 430m model with "The capital of France is" it continues with "first of the, the city of Paris". I checked that the tokenizer is working properly, so I think it's something to do with the inference / context-handling code.

Some other random notes:

• The models were twice their size when porting to ONNX - e.g. 169m model goes from 339MB to 679MB. I quantized down to 171MB, but that makes inference half the speed (~5 tokens/sec for quantized versus ~13 tokens/sec for original). I'm guessing the non-quantized version have been converted from f16 to f32, hence the size doubling? The demo includes both the normal and quantized versions.

• @BlinkDL Yes, I got TypeError: cannot resolve operator 'Max' with opsets: ai.onnx v13 when trying to use the WebGL backend. How would I go about removing RWKV_HEAD_QK and RWKV-ffnPre? I made a conversion notebook here: https://colab.research.google.com/github/josephrocca/rwkv-v4-web/blob/main/RWKV_v4_ONNX_conversion.ipynb Is it as simple as adding a few commands to that, or is there more work involved?

• The WebGL backend doesn't work with quantized models. It gives this error: TypeError: cannot resolve operator 'DequantizeLinear' with opsets: ai.onnx v13, com.microsoft.experimental v1, ai.onnx.preview.training v1, ai.onnx.training v1, com.ms.internal.nhwc v17, org.pytorch.aten v1, com.microsoft.nchwc v1, ai.onnx.ml v3, com.microsoft v1

• I used a very overkill approach to getting the tokenizer working... https://github.com/josephrocca/tokenizers-pyodide I haven't looked into how different the tokenizer is from gpt 2/3, but if it's similar, then I guess it shouldn't be too hard to make an edited version of this https://github.com/josephrocca/gpt-2-3-tokenizer ?

[AXKuhta]

For example, if you prompt the 430m model with "The capital of France is" it continues with "first of the, the city of Paris"

That seems familiar!

The => first
The capital => of
The capital of => the
The capital of France => ,
The capital of France is => the

It looks like you display the outputs during the prompt-feeding stage, which happens one token at a time.

That should fix it:

         let token = greedySampling(results.x.data);

         if (promptTokens.length == 0) {
+          if(streamingCallback) streamingCallback(token);
           ctx.push( token );
         } else {
           ctx.push( promptTokens.shift() );
         }
-
-        if(streamingCallback) streamingCallback(token);

         feeds.xx_att = results.xx_att_r;
         feeds.aa_att = results.aa_att_r;

[AXKuhta]

@josephrocca I had to host the demo locally because huggingface keeps terminating the model downloads for some reason, but otherwise I can confirm that it works on my machine. Good job with getting the tokenizer and the quantization working!

I'm guessing the non-quantized version have been converted from f16 to f32, hence the size doubling?

Yeah, that's what's happening. RWKV-v4 is bf16 which can't be losslessly converted to fp16, so fp32 is the next best option. The fp32-converted model also compresses really well since half the bytes in it are zero.

[BlinkDL]

• @BlinkDL Yes, I got TypeError: cannot resolve operator 'Max' with opsets: ai.onnx v13 when trying to use the WebGL backend. How would I go about removing RWKV_HEAD_QK and RWKV-ffnPre? I made a conversion notebook here: https://colab.research.google.com/github/josephrocca/rwkv-v4-web/blob/main/RWKV_v4_ONNX_conversion.ipynb Is it as simple as adding a few commands to that, or is there more work involved?

take a look at src/model_run.py.
for the pile model, self.model_type == 'RWKV' and RWKV_HEAD_QK_DIM = 0 so you can remove some useless code.

moreover, use https://github.com/daquexian/onnx-simplifier to optimizer the onnx model

[BlinkDL]

And the onnx version might work for AMD & Intel gpus. The DirectML backend supports them (on win10).

I tried that for RWKV-1.

[BlinkDL]

Yeah, that's what's happening. RWKV-v4 is bf16 which can't be losslessly converted to fp16, so fp32 is the next best option. The fp32-converted model also compresses really well since half the bytes in it are zero.

You can loseless "transform" bf16 to fp16, and the idea is to use the same raw binary value. The float value will be totally different, but you can do an inverse transform in JS to loselessly recover the original bf16.

[josephrocca]

@AXKuhta Could have sworn I replied here earlier, sorry - apparently I didn't click send. I fixed the demo according to your comment soon after you posted it - thanks for your help!! Strange that huggingface is terminating the download for you... 🤔

@BlinkDL Thanks for the tips! I'll look into the stuff you've mentioned.

[BlinkDL]

Hi, really exciting project! I'm wondering if you've published the model conversion script that you used to create the js_models files from the .pth model file? It would be awesome to see how the larger and newer models like RWKV-4 169m/430m perform in the browser! I think the inference speed of RWKV opens up many new possibilities for language models on the web.

The python code for RWKV-2 weight conversion to .bin (for tf.js):

w = torch.load(MODEL_NAME + '.pth')
for x in w.keys():
	if 'copy_mask' in x: # this is for headQK which is not used in pile models
		continue
	print(x, w[x].shape)
	
	# we are doing some pre-computations here. change them to match RWKV-4. or you can just skip all of them and do everything in js first.
	if '.time_' in x: 
		w[x] = w[x].squeeze()
	if '.time_decay' in x:
		w[x] = torch.exp(-torch.exp(w[x]))
	if '.time_first' in x:
		w[x] = torch.exp(w[x])
	
	w[x].numpy().tofile(f'20220425/{x}.bin')

You can gradually port it to RWKV-4 by matching the outputs for each layer.

The Chinese RWKV-2 has a better UI: https://github.com/BlinkDL/AI-Writer/blob/main/docs/index.html

The English RWKV-2: https://github.com/BlinkDL/AI-Writer/tree/main/docs/eng

[BlinkDL]

@AXKuhta Could have sworn I replied here earlier, sorry - apparently I didn't click send. I fixed the demo according to your comment soon after you posted it

Add top-p top-k and temperature and then it's very usable :)

[AXKuhta]

It looks like the webgl backend has a lot of limitations. I did some testing by stripping out different parts of the model in order to see if I can get anything at all to work on the webgl backend. I think I got like four different error messages with different combinations. The bottom line is that I can't even get a matmul to work.

https://github.com/AXKuhta/RWKV-LM/blob/matmul/RWKV-v4/matmul_test.py
https://github.com/AXKuhta/RWKV-LM/blob/matmul/RWKV-v4/index.html

It does work on the wasm backend!

EDIT: It actually works on webgl if you do this: AXKuhta@75ad160

[AXKuhta]

I have been able to force the full model to run on webgl, but it doesn't produce anything coherent, so something's still broken:

https://github.com/AXKuhta/RWKV-LM/tree/onnx_webgl

@BlinkDL The "cannot resolve operator 'Max' with opsets: ai.onnx v13" error was caused by torch.maximum(pp, ww) and I was able to suppress it by using torch.max(torch.stack([pp, ww]), 0).values instead. I also had to add a bunch of .view([768,1]) around matmul operations and then fix layer_norm() from producing NaNs. Now it looks like self.FF() always produces zeroes on webgl, but I'm not sure yet as to why nope, self.FF() does produce something.

[BlinkDL]

I have been able to force the full model to run on webgl, but it doesn't produce anything coherent, so something's still broken:

https://github.com/AXKuhta/RWKV-LM/tree/onnx_webgl

@BlinkDL The "cannot resolve operator 'Max' with opsets: ai.onnx v13" error was caused by torch.maximum(pp, ww) and I was able to suppress it by using torch.max(torch.stack([pp, ww]), 0).values instead. I also had to add a bunch of .view([768,1]) around matmul operations and then fix layer_norm() from producing NaNs. Now it looks like self.FF() always produces zeroes on webgl, but I'm not sure yet as to why nope, self.FF() does produce something.

That's great. Could you check whether https://github.com/daquexian/onnx-simplifier can help? Use https://github.com/lutzroeder/netron to visualize models.

And then you can print() the outputs of interesting layers to find the culprit... gradually matching the results of webgl vs wasm.

[AXKuhta]

@BlinkDL After some painstaking debugging I got it to produce coherent output on webgl. The fix was really bizarre: add + 0.0 in a bunch of places. Some nodes on the ONNX graph that follow matmul+reshape operations kept getting bugged inputs that looked like a single value across all 768 elements. Performing +0.0 with the bugged input fixes it.

Here's the changes: https://github.com/AXKuhta/RWKV-LM/commits/onnx_webgl

Could you check whether https://github.com/daquexian/onnx-simplifier can help?

I did try onnx-simplifier with RWKV-3, but it didn't find much to simplify. The graph was almost unchanged. I will retest with RWKV-4 though.