AI/text-generation-webui
1
0
Fork 0
mirror of https://github.com/oobabooga/text-generation-webui.git synced 2024-10-01 01:26:03 -04:00
Code Issues Packages Projects Releases Wiki Activity

auto-assign gpu memory with --auto-devices alone

Browse Source
This commit is contained in:
oobabooga 2023-02-10 16:36:06 -03:00
parent 76d3d7ddb3
commit 316e07f06a
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
server.py
View File

@ -130,7 +130,7 @@ def load_model(model_name):
if args.gpu_memory: if args.gpu_memory:
params.append(f"max_memory={{0: '{args.gpu_memory or '99'}GiB', 'cpu': '{args.cpu_memory or '99'}GiB'}}") params.append(f"max_memory={{0: '{args.gpu_memory or '99'}GiB', 'cpu': '{args.cpu_memory or '99'}GiB'}}")
elif (args.gpu_memory or args.cpu_memory) and not args.load_in_8bit: elif not args.load_in_8bit:
total_mem = (torch.cuda.get_device_properties(0).total_memory/(1024*1024)) total_mem = (torch.cuda.get_device_properties(0).total_memory/(1024*1024))
suggestion = round((total_mem-1000)/1000)*1000 suggestion = round((total_mem-1000)/1000)*1000
if total_mem-suggestion < 800: if total_mem-suggestion < 800: