mmarch/text-generation-webui
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Add llama.cpp GPU offload option (#2060)

Browse source
This commit is contained in:
AlphaAtlas 2023-05-14 21:58:11 -04:00 • committed by GitHub
parent eee986348c
commit 071f0776ad
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 23 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

3
modules/llamacpp_model.py
View file

@ -27,7 +27,8 @@ class LlamaCppModel:
'n_threads': shared.args.threads or None,
'n_batch': shared.args.n_batch,
'use_mmap': not shared.args.no_mmap,
'use_mlock': shared.args.mlock
'use_mlock': shared.args.mlock,
'n_gpu_layers': shared.args.n_gpu_layers
}
self.model = Llama(**params)
self.model.set_cache(LlamaCache)

1
modules/shared.py
View file

@ -123,6 +123,7 @@ parser.add_argument('--threads', type=int, default=0, help='Number of threads to
parser.add_argument('--n_batch', type=int, default=512, help='Maximum number of prompt tokens to batch together when calling llama_eval.')
parser.add_argument('--no-mmap', action='store_true', help='Prevent mmap from being used.')
parser.add_argument('--mlock', action='store_true', help='Force the system to keep the model in RAM.')
parser.add_argument('--n-gpu-layers', type=int, default=0, help='Number of layers to offload to the GPU.')
# 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.')