Chat demo does not work on Android because of maxStorageBufferBindingSize

[beaufortfrancois]

On my Pixel 7 Android device, the maxStorageBufferBindingSize limit is only 128. Sadly https://webllm.mlc.ai/#chat-demo requires 1024 for the Llama-2-7b-chat-hf-q4f16_1 model as seen below when run on Chrome for Android.

As it is possible with an Android app (https://llm.mlc.ai/#android), would it be possible to lower this limit so that more Android devices can access and play with https://webllm.mlc.ai/#chat-demo? Note that WebGPU support is coming to Android soon. See https://groups.google.com/a/chromium.org/g/blink-dev/c/YFWuDlCKTP4/m/97C4LCBUBgAJ

[beaufortfrancois]

It seems related to apache/tvm#14640

[tqchen]

thanks @beaufortfrancois for bringing it up!

Unfortunately this is a limit that we would need help from Chrome side to lift. This is mainly because the model itself do require buffers beyond that 128MB(we need bigger storages of some of the parameters and kv cache).

In our native android demo, we relied on OpenCL that does not have this limit. This being said, llama2 on android is indeed a super heavy lifting as of now and requires the latest snapdragon device to function (eating about 5-6G of vRAM).

RedPajama 3B model might be a good candidate if we want to try to see the limit there

[beaufortfrancois]

How can the Chrome team assist you in addressing this limitation? The maxStorageBufferBindingSize WebGPU limit is actually derived from Vulkan's maxStorageBufferRange limit.

[tqchen]

Thanks for the note! let us look a bit into it and see if it is possible to get a variant of small model that fit into this limit.

In the meantime, I feel that the vulkan driver should be able to lift some of the limit, for reference, we had some opencl results https://blog.mlc.ai/2023/08/09/GPU-Accelerated-LLM-on-Orange-Pi based on Mali so the GPU itself should able to handle bigger range, but of course that would require changes to the vulkan driver end

[CharlieFRuan]

Hi @beaufortfrancois, the demo now has three models (ones with -1k suffix) that should be able to run with 128MB: https://webllm.mlc.ai/

[beaufortfrancois]

First of all, thanks for providing those models!

I've just tried the Llama-2-7b-chat-hf-q4f16_1-1k model in Chrome Canary for Android where WebGPU is enabled by default and as you can see below, the never-ending answer seems wrong when asking "Which color is the dress?". Is this an issue in the model or in the implementation?

Note that it works great on "Apple M1 Pro" device.

screen-20231113-091318.mp4

[tqchen]

that model depends on shader-f16 feature, which only existed in chrom canary (and not chrome stable AFAIK), so not sure if it works on android. If it is possible to get the console output, it will also help us see what is going on

Do you mind trying out RedPajama-INCITE-Chat-3B-v1-q32-1k?

[beaufortfrancois]

that model depends on shader-f16 feature, which only existed in chrom canary (and not chrome stable AFAIK), so not sure if it works on android. If it is possible to get the console output, it will also help us see what is going on

shader-f16 is shipping in Chrome 120. See https://groups.google.com/a/chromium.org/g/blink-dev/c/AsKn-UwMYAE/m/4FKB-x_QAQAJ

I'm unable to get logs as Chrome is now crashing with this model. Sorry ;( I've tried several times.

Do you mind trying out RedPajama-INCITE-Chat-3B-v1-q32-1k?

It works great with RedPajama-INCITE-Chat-3B-v1-q4f32_1-1k.

[tqchen]

It is possible that the model crashes in Llama, we can hit VRAM limit in Llama2 models(that goes beyond 4GB) when using 4 bit quantization, and in iOS we had to rely on 3bit quantization to stay in memory budget. On android we can sometimes get Llama2 4bit model working but that also requires to go beyond 4GB VRAM limit and the depending on the system limit and phone's capability it can crash.

So likely 3B models are good choice to demo LLM on android/webgpu, or if later we have 2bit model variants. Glad that RedPajama-INCITE-Chat-3B-v1-q4f32_1-1k works, perhaps you can also try RedPajama-INCITE-Chat-3B-v1-q4f16_1-1k to test f16 support

[tqchen]

Glad that 3B model works, this is the first running example of webgpu native LLM on mobile phone AFAIK, thank you @beaufortfrancois for pushing this. Love to share this with broader community and let us know how we can help

[beaufortfrancois]

It is possible that the model crashes in Llama, we can hit VRAM limit in Llama2 models(that goes beyond 4GB) when using 4 bit quantization, and in iOS we had to rely on 3bit quantization to stay in memory budget. On android we can sometimes get Llama2 4bit model working but that also requires to go beyond 4GB VRAM limit and the depending on the system limit and phone's capability it can crash.

🥳 I was finally able to get logs on my Pixel 7 Android device!

[FATAL] /ssd1/cfruan/packages/tvm-unity/src/runtime/memory/memory_manager.cc:108:
InternalError: Check failed: (offset + needed_size <= this->buffer.size) is false: 
storage allocation failure, attempted to allocate 8396800 at offset 0 in region that is 8388608bytes
@ worker.360c8a9c.js:669

Uncaught (in promise) DOMException: Failed to execute 'mapAsync' on 'GPUBuffer': Device is lost

Hopefully this will help.

So likely 3B models are good choice to demo LLM on android/webgpu, or if later we have 2bit model variants. Glad that RedPajama-INCITE-Chat-3B-v1-q4f32_1-1k works, perhaps you can also try RedPajama-INCITE-Chat-3B-v1-q4f16_1-1k to test f16 support

Glad that 3B model works, this is the first running example of webgpu native LLM on mobile phone AFAIK, thank you @beaufortfrancois for pushing this. Love to share this with broader community and let us know how we can help

I'm also really happy 😄 to see WebLLM running on Android without any flag ;)

[tqchen]

The error log indeed indicate possible OOM ("gpu get lost") because the model is too VRAM demanding