Enabling CUDA on Docker

[franco-giordano]

Hi, I'm having trouble enabling CUDA-powered inference on my machine. What I've tried:

• Use default docker compose -> no luck
• Git clone repo, enable Nvidia driver in docker compose and build locally with a CUDA-based image in Dockerfile (FROM nvidia/cuda:12.2.0-devel-ubuntu22.04) -> pytorch can't detect CUDA GPU
• Same as above, but using an older image (FROM nvidia/cuda:11.7.1-devel-ubuntu22.04) -> pytorch detects CUDA, but Khoj loads LLM to CPU (Loaded chat model to CPU. utils.py:32 in logs, CPU usage spikes)

Any ideas? From previous discussions, issues, and docs I get the impression that it should work, but I'm not sure. Some specs:

• Default Offline model: mistral-7b-instruct
• GPU: GTX 1660 SUPER (6GB VRAM)
• nvidia-smi output:

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.125.06   Driver Version: 525.125.06   CUDA Version: 12.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA GeForce ...  On   | 00000000:01:00.0  On |                  N/A |
| 38%   35C    P8    12W / 125W |    236MiB /  6144MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A       852      G   /usr/lib/xorg/Xorg                 81MiB |
|    0   N/A  N/A      1018    C+G   ...ome-remote-desktop-daemon      107MiB |
|    0   N/A  N/A      1052      G   /usr/bin/gnome-shell               26MiB |
|    0   N/A  N/A   2751199      G   /usr/bin/nautilus                  15MiB |
+-----------------------------------------------------------------------------+

• Pytorch inside the container:

$ docker exec -it khoj-server-1 bash
root@3970858b56e9:/app# python3
Python 3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch.cuda.is_available()
True

[debanjum]

Hey @franco-giordano, thanks for the detailed report:

1. You'd need a Vulkan supported GPU. GTX 1660 SUPER is supported
2. Mistral practically needs at least 8Gb VRAM (ideally 16Gb) to load/run. This seems like the most likely reason Khoj is loading Mistral onto CPU

[franco-giordano]

Thanks for the quick response! That must be it then. This raises a few questions though:

• does Khoj only support GPT4ALL's models? I see that there is only 15 models. Or does it support any from HuggingFace? (i.e. I've tried before using mistral-7b-v0.1.Q4_K_M.gguf but got a 404 error)
• also, any tips on where to find how much VRAM is required for these models? I'm guessing its simply the RAM, not sure though
• Assuming I find a suitable model, do I have to maintain my Dockerfile/docker-compose tweaks? Or should it work with latest?

Thanks again!

[debanjum]

• I haven't tested it but you maybe able to use any model in gguf format. You'll need to manually download a chat model in gguf format into the ~/.cache/gpt4all/ folder

• The VRAM is the same as the RAM required by the model. When model can fit into VRAM/GPU it will be loaded there instead of into RAM for faster responses

• What tweaks have you made to your docker-compose.yml ? You can pull the latest Khoj docker image, if you want to persist some settings/changes you can use docker volumes for that

[DeutscheGabanna]

Can you share your docker-compose.yaml so that I can see how you integrated CUDA into the script?