Mirrors/text-generation-webui
1
0
Fork 0
mirror of https://github.com/oobabooga/text-generation-webui.git synced 2024-09-20 10:35:10 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

Try to leave at least 1GiB free to prevent oom errors

Browse source
This commit is contained in:
oobabooga 2023-01-31 20:47:05 -03:00
parent bb77f20a6c
commit 8ef89df746
1 changed files with 10 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
server.py
View file

@ -97,10 +97,16 @@ def load_model(model_name):
settings.append("device_map='auto'")
settings.append("load_in_8bit=True" if args.load_in_8bit else "torch_dtype=torch.float16")
if args.gpu_memory and args.cpu_memory:
settings.append(f"max_memory={{0: '{args.gpu_memory}GiB', 'cpu': '{args.cpu_memory}GiB'}}")
elif args.gpu_memory:
settings.append(f"max_memory={{0: '{args.gpu_memory}GiB', 'cpu': '99GiB'}}")
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:
total_mem = (torch.cuda.get_device_properties(0).total_memory/(1024*1024))
suggestion = round((total_mem-1000)/1000)*1000
if total_mem-suggestion < 800:
suggestion -= 1000
suggestion = int(round(suggestion/1000))
print(f"\033[1;32;1mAuto-assiging --gpu-memory {suggestion} for your GPU to try to prevent out-of-memory errors.\nYou can manually set other values.\033[0;37;0m")
settings.append(f"max_memory={{0: '{suggestion}GiB', 'cpu': '{args.cpu_memory or '99'}GiB'}}")
if args.disk:
settings.append(f"offload_folder='{args.disk_cache_dir or 'cache'}'")