Scores are stored in a 32-bit NumPy array even when K and V are quantized

Open Beginner friendly
#1,732 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Assessment

Difficulty
2/5
Estimated time
1-3 hours
Newbie friendliness
65/100
Issue type
Bug
Clarity
Clearly specified
Activity status
Stale
Tech stack
numpy, python

Research direction

The issue points to line 460 in llama_cpp/llama.py where a NumPy array is created with dtype=np.single. The fix is to change the dtype to match the quantization of the K and V tensors (e.g., np.half for 16-bit). First, examine the surrounding code to understand how the scores array is used and what dtype the quantized tensors have. Then, modify the dtype accordingly and test by running the server command with a large context to ensure no memory error occurs.

Written by the indexing model from the issue text.

Description

Prerequisites

Please answer the following questions for yourself before submitting an issue.

  • I am running the latest code. Development is very rapid so there are no tagged versions as of now.
  • I carefully followed the README.md.
  • I searched using keywords relevant to my issue to make sure that I am creating a new issue that is not already open (or closed).
  • I reviewed the Discussions, and have a new bug or useful enhancement to share.

Expected Behavior

It should load the scores into an array with the appropriate data type

Current Behavior

Instead, it loads them into 32 bit ndarray

Environment and Context

Windows 10
Python 3.11.9
Latest CUDA 12.1 wheel as of now
RTX 3090
RTX 3070 (Hidden to application in this test)
32GB RAM

Failure Information

Traceback (most recent call last):
  File "<frozen runpy>", line 198, in _run_module_as_main
  File "<frozen runpy>", line 88, in _run_code
  File "C:\Users\Ethan\miniconda3\envs\torch\Lib\site-packages\llama_cpp\server\__main__.py", line 100, in <module>
    main()
    app = create_app(
          ^^^^^^^^^^^
  File "C:\Users\Ethan\miniconda3\envs\torch\Lib\site-packages\llama_cpp\server\app.py", line 150, in create_app
    set_llama_proxy(model_settings=model_settings)
  File "C:\Users\Ethan\miniconda3\envs\torch\Lib\site-packages\llama_cpp\server\app.py", line 70, in set_llama_proxy
    _llama_proxy = LlamaProxy(models=model_settings)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\Ethan\miniconda3\envs\torch\Lib\site-packages\llama_cpp\server\model.py", line 31, in __init__
    self._current_model = self.load_llama_from_model_settings(
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\Ethan\miniconda3\envs\torch\Lib\site-packages\llama_cpp\server\model.py", line 236, in load_llama_from_model_settings
    _model = create_fn(
             ^^^^^^^^^^
  File "C:\Users\Ethan\miniconda3\envs\torch\Lib\site-packages\llama_cpp\llama.py", line 460, in __init__
    self.scores: npt.NDArray[np.single] = np.ndarray(
                                          ^^^^^^^^^^^
numpy.core._exceptions._ArrayMemoryError: Unable to allocate 113. GiB for an array with shape (200000, 151552) and data type float32

Steps to Reproduce

Use the openai server with this command (or similar):
python -m llama_cpp.server --model glm-4-9b-chat-1m-Q4_0.gguf --flash_attn true --type_k 6 --type_v 6 --n_gpu_layers -1 --n_ctx 200000

Solution

The issue is on this line:
https://github.com/abetlen/llama-cpp-python/blob/main/llama_cpp/llama.py#L460C9-L462C10
I set it to dtype=np.half and it worked, but my kv quants are still at 6bit so I think it could go lower.

Dominant language
Python
Stars
10.6k
Forks
1.5k
Avg merge
23m
Merged PRs (30d)
1

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from abetlen/llama-cpp-python

All issues in abetlen/llama-cpp-python

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.