zhengjl/text-generation-webui
Public Access
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Merge branch 'oobabooga:main' into ds

Browse Source
This commit is contained in:
81300
2023-02-01 20:50:51 +02:00
committed by GitHub
parent a6f4760772 6b13816c47
commit 248ec4fa21
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
server.py
View File

@@ -200,7 +200,7 @@ def load_model(model_name):
# Custom
else:
command = "AutoModelForCausalLM.from_pretrained"
settings = []
settings = ["low_cpu_mem_usage=True"]
if args.cpu:
settings.append("low_cpu_mem_usage=True")
@@ -211,7 +211,7 @@ def load_model(model_name):
if args.gpu_memory:
settings.append(f"max_memory={{0: '{args.gpu_memory or '99'}GiB', 'cpu': '{args.cpu_memory or '99'}GiB'}}")
elif not args.load_in_8bit:
elif (args.gpu_memory or args.cpu_memory) and not args.load_in_8bit:
total_mem = (torch.cuda.get_device_properties(0).total_memory/(1024*1024))
suggestion = round((total_mem-1000)/1000)*1000
if total_mem-suggestion < 800: