Repository: oobabooga/text-generation-webui
Pull request: Add the --cpu option for llama.cpp to prevent CUDA from being used
Fixes #3412. 

The new `llama_cpp_lib()` function can later be used to support the AMD version of llama-cpp-python while keeping the rest of the code unchanged.

cc @jllllll 

Problem statement:
Please make an option to *NOT* use the built-in llama.cpp + cuBLAS support
**Description**

Until a recent update I was building` llama-cpp-python` for `text-generation-webui` myself, to make sure cuBLAS is supported (for my RTX 3060). But since my CPU does not support certain modern instructions (AVX2, F16C and FMA) I have to make sure to uncheck those flags when building with `cmake`. But now I cannot use that anymore, since it loads the llama_cpp_cuda module by default, which tries to use instructions that my CPU does not understand and therefore crashes with a core dump.

llamacpp_model.py at the moment:

```
if torch.cuda.is_available() and not torch.version.hip:
    try:
        from llama_cpp_cuda import Llama, LlamaCache, LogitsProcessorList
    except:
        from llama_cpp import Llama, LlamaCache, LogitsProcessorList
else:
    from llama_cpp import Llama, LlamaCache, LogitsProcessorList
```

**Additional Context**

My current workaround is to make sure that llama_cpp is loaded everytime, even when CUDA is detected. But that's a horrible way to make sure I run my stuff.

Hints from issue discussion:
Currently, the easiest workaround for this is to uninstall `llama-cpp-python-cuda`. If it is missing, it will default to the
`llama-cpp-python` package, allowing you to use your custom build. A more reliable solution would be to add a check for the `--cpu` flag like so:
```python
if torch.cuda.is_available() and not torch.version.hip and not shared.args.cpu:
    try:
        from llama_cpp_cuda import Llama, LlamaCache, LogitsProcessorList
    except:
        from llama_cpp import Llama, LlamaCache, LogitsProcessorList
else:
    from llama_cpp import Llama, LlamaCache, LogitsProcessorList
```
Either that, or use `py-cpuinfo` to check what the CPU supports:
```python
from cpuinfo import get_cpu_info

if torch.cuda.is_available() and not torch.version.hip and all(cpuflag in get_cpu_info()['flags'] for cpuflag in ['avx2','f16c','fma']):
    try:
        from llama_cpp_cuda import Llama, LlamaCache, LogitsProcessorList
    except:
        from llama_cpp import Llama, LlamaCache, LogitsProcessorList
else:
    from llama_cpp import Llama, LlamaCache, LogitsProcessorList
```
The `py-cpuinfo` option will introduce a noticeable delay in loading llama-cpp-python.

Could also just add a new option for not using llama-cpp-python-cuda as you suggest.
A third option could be to import both `llama_cpp` and `llama_cpp_cuda` if `torch.cuda.is_available() and not torch.version.hip`, and then choose which `Llama, LlamaCache, LogitsProcessorList` to use at runtime depending on `shared.args.cpu`. This way, the user can switch between one `llama_cpp` or the other on the fly.

Also, this makes me wonder about AMD. ExLlama, AutoGPTQ, and GPTQ-for-LLaMa are all based on CUDA, so I guess they don't work for AMD. Transformers is based on pytorch, which will work on Linux with ROCm but not Windows. So llama-cpp-python with GPU acceleration is the best option. Do you see a route to get llama-cpp-python to work out of the box on Windows and Linux for AMD @jllllll?
@oobabooga

There is really only one way to have AMD GPU support for both Windows and Linux: Build llama-cpp-python for CLBlast support.

This has the potential of providing worse performance than just CPU-only due to the need to "tune" the CLBlast kernels for a given GPU. It comes pre-tuned for these GPUs:
https://github.com/CNugteren/CLBlast/blob/master/doc/tuning.md
It's best performance is still noticeably slower than using hipBLAS, which is only available on Linux for a limited number of GPUs.
However, it is still better than nothing for the GPUs that it supports, especially on Windows where this is the only option.
The main issue that I've encountered with this is that the CLBlast library needs to be in a folder listed in `PATH` to be loaded on Windows. This can be automated in the one-click-installer, but will not work out of the box with a manual installation without manually installing CLBlast libs.

There is also this AMD port of llama.cpp's cuBLAS kernels: https://github.com/SlyEcho/llama.cpp/tree/hipblas
This will only work for Linux and doesn't support WSL.
Using this may be tricky as it requires that the fork commit used is based on a version of llama.cpp that is compatible with the version of llama-cpp-python that is being built. Another issue is that it will only work on the GPUs that Pytorch supports due to the low compatibility of hipBLAS.
Finally figured out why loading CLBlast on Windows is so difficult. A single line of code in llama-cpp-python's `llama_cpp.py` file is overriding the search paths for Windows and preventing the bundled CLBlast libs from being located:
```python
cdll_args["winmode"] = 0
```
Not sure why this was added back in May. Removing it does not break any functionality as far as I can tell. I created an issue in the llama-cpp-python repo to see if it can be removed or if an alternative solution can be implemented: https://github.com/abetlen/llama-cpp-python/issues/563
If this is fixed, including llama-cpp-python CLBlast support will be trivial.
> However, it is still better than nothing for the GPUs that it supports

That's what I think as well. Most AMD users bought their cards for Windows gaming, so supporting those users would be a huge progress already.

> A single line of code in llama-cpp-python's llama_cpp.py file is overriding the search paths

That's a really nice catch! Let's wait for it to get resolved then.

Commit message: README updates

Patch to review:
diff --git a/README.md b/README.md
index 073a841dd2..6ec84ba247 100644
--- a/README.md
+++ b/README.md
@@ -249,8 +249,9 @@ Optionally, you can use the following command-line flags:
 | `--n-gpu-layers N_GPU_LAYERS` | Number of layers to offload to the GPU. Only works if llama-cpp-python was compiled with BLAS. Set this to 1000000000 to offload all layers to the GPU. |
 | `--n_ctx N_CTX` | Size of the prompt context. |
 | `--llama_cpp_seed SEED` | Seed for llama-cpp models. Default 0 (random). |
-| `--n_gqa N_GQA`         | grouped-query attention. Must be 8 for llama2 70b. |
-| `--rms_norm_eps RMS_NORM_EPS`  | Must be 1e-5 for llama2 70b. |
+| `--n_gqa N_GQA`         | grouped-query attention. Must be 8 for llama-2 70b. |
+| `--rms_norm_eps RMS_NORM_EPS`  | 5e-6 is a good value for llama-2 models. |
+| `--cpu`                        | Use the CPU version of llama-cpp-python instead of the GPU-accelerated version. |
 
 #### AutoGPTQ
 
diff --git a/modules/llamacpp_hf.py b/modules/llamacpp_hf.py
index 349a578209..465c81501d 100644
--- a/modules/llamacpp_hf.py
+++ b/modules/llamacpp_hf.py
@@ -12,11 +12,20 @@
 
 if torch.cuda.is_available() and not torch.version.hip:
     try:
-        from llama_cpp_cuda import Llama
+        import llama_cpp_cuda
     except:
-        from llama_cpp import Llama
+        llama_cpp_cuda = None
+        import llama_cpp
 else:
-    from llama_cpp import Llama
+    llama_cpp_cuda = None
+    import llama_cpp
+
+
+def llama_cpp_lib():
+    if shared.args.cpu or llama_cpp_cuda is None:
+        return llama_cpp
+    else:
+        return llama_cpp_cuda
 
 
 class LlamacppHF(PreTrainedModel):
@@ -111,5 +120,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
             'logits_all': True,
         }
 
+        Llama = llama_cpp_lib().Llama
         model = Llama(**params)
+
         return LlamacppHF(model)
diff --git a/modules/llamacpp_model.py b/modules/llamacpp_model.py
index 0f9c3470ca..8a397beb29 100644
--- a/modules/llamacpp_model.py
+++ b/modules/llamacpp_model.py
@@ -9,11 +9,20 @@
 
 if torch.cuda.is_available() and not torch.version.hip:
     try:
-        from llama_cpp_cuda import Llama, LlamaCache, LogitsProcessorList
+        import llama_cpp_cuda
     except:
-        from llama_cpp import Llama, LlamaCache, LogitsProcessorList
+        llama_cpp_cuda = None
+        import llama_cpp
 else:
-    from llama_cpp import Llama, LlamaCache, LogitsProcessorList
+    llama_cpp_cuda = None
+    import llama_cpp
+
+
+def llama_cpp_lib():
+    if shared.args.cpu or llama_cpp_cuda is None:
+        return llama_cpp
+    else:
+        return llama_cpp_cuda
 
 
 def ban_eos_logits_processor(eos_token, input_ids, logits):
@@ -30,6 +39,10 @@ def __del__(self):
 
     @classmethod
     def from_pretrained(self, path):
+
+        Llama = llama_cpp_lib().Llama
+        LlamaCache = llama_cpp_lib().LlamaCache
+
         result = self()
         cache_capacity = 0
         if shared.args.cache_capacity is not None:
@@ -74,6 +87,9 @@ def decode(self, tokens):
         return self.model.detokenize(tokens)
 
     def generate(self, prompt, state, callback=None):
+
+        LogitsProcessorList = llama_cpp_lib().LogitsProcessorList
+
         prompt = prompt if type(prompt) is str else prompt.decode()
         completion_chunks = self.model.create_completion(
             prompt=prompt,
diff --git a/modules/loaders.py b/modules/loaders.py
index 68b48204ed..aa1afcb854 100644
--- a/modules/loaders.py
+++ b/modules/loaders.py
@@ -41,6 +41,7 @@
         'llama_cpp_seed',
         'compress_pos_emb',
         'alpha_value',
+        'cpu',
     ],
     'llamacpp_HF': [
         'n_ctx',
@@ -55,6 +56,7 @@
         'llama_cpp_seed',
         'compress_pos_emb',
         'alpha_value',
+        'cpu',
         'llamacpp_HF_info',
     ],
     'Transformers': [
diff --git a/modules/shared.py b/modules/shared.py
index bac3fa8cc6..fc9ba3cf19 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -132,8 +132,8 @@ def str2bool(v):
 parser.add_argument('--n-gpu-layers', type=int, default=0, help='Number of layers to offload to the GPU.')
 parser.add_argument('--n_ctx', type=int, default=2048, help='Size of the prompt context.')
 parser.add_argument('--llama_cpp_seed', type=int, default=0, help='Seed for llama-cpp models. Default 0 (random)')
-parser.add_argument('--n_gqa', type=int, default=0, help='grouped-query attention. Must be 8 for llama2 70b.')
-parser.add_argument('--rms_norm_eps', type=float, default=0, help='Must be 1e-5 for llama2 70b.')
+parser.add_argument('--n_gqa', type=int, default=0, help='grouped-query attention. Must be 8 for llama-2 70b.')
+parser.add_argument('--rms_norm_eps', type=float, default=0, help='5e-6 is a good value for llama-2 models.')
 
 # GPTQ
 parser.add_argument('--wbits', type=int, default=0, help='Load a pre-quantized model with specified precision in bits. 2, 3, 4 and 8 are supported.')


Code review comments:
diff_hunk: @@ -12,11 +12,20 @@
 
 if torch.cuda.is_available() and not torch.version.hip:
     try:
-        from llama_cpp_cuda import Llama
+        import llama_cpp_cuda
     except:
-        from llama_cpp import Llama
+        llama_cpp_cuda = None
+        import llama_cpp
path: modules/llamacpp_hf.py
text: @user1:
If you remove the indent from this line, then you can freely swap between llama_cpp_cuda and llama_cpp, which I believe is what you said that you wanted. Same with `llamacpp_model.py`
```python
    try:
        from llama_cpp_cuda import Llama
        import llama_cpp_cuda
    except:
        from llama_cpp import Llama
        llama_cpp_cuda = None
    import llama_cpp
```

@author:
Oops, my bad, I missed that. It should be fixed now.