mirror of
https://github.com/oobabooga/text-generation-webui.git
synced 2024-10-01 01:26:03 -04:00
get_max_memory_dict() was not properly formatting shared.args.cpu_memory Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
This commit is contained in:
parent
8a9d5a0cea
commit
025da386a0
@ -381,15 +381,13 @@ def RWKV_loader(model_name):
|
||||
|
||||
|
||||
def get_max_memory_dict():
|
||||
max_memory = {}
|
||||
max_cpu_memory = shared.args.cpu_memory.strip() if shared.args.cpu_memory is not None else '99GiB'
|
||||
max_memory = {'cpu': f'{max_cpu_memory}GiB' if not re.match('.*ib$', max_cpu_memory.lower()) else max_cpu_memory}
|
||||
if shared.args.gpu_memory:
|
||||
memory_map = list(map(lambda x: x.strip(), shared.args.gpu_memory))
|
||||
for i in range(len(memory_map)):
|
||||
max_memory[i] = f'{memory_map[i]}GiB' if not re.match('.*ib$', memory_map[i].lower()) else memory_map[i]
|
||||
|
||||
max_cpu_memory = shared.args.cpu_memory.strip() if shared.args.cpu_memory is not None else '99GiB'
|
||||
max_memory['cpu'] = f'{max_cpu_memory}GiB' if not re.match('.*ib$', max_cpu_memory.lower()) else max_cpu_memory
|
||||
|
||||
# If --auto-devices is provided standalone, try to get a reasonable value
|
||||
# for the maximum memory of device :0
|
||||
elif shared.args.auto_devices:
|
||||
@ -403,7 +401,7 @@ def get_max_memory_dict():
|
||||
|
||||
suggestion = int(round(suggestion / 1000))
|
||||
logger.warning(f"Auto-assiging --gpu-memory {suggestion} for your GPU to try to prevent out-of-memory errors. You can manually set other values.")
|
||||
max_memory = {0: f'{suggestion}GiB', 'cpu': f'{shared.args.cpu_memory or 99}GiB'}
|
||||
max_memory[0] = f'{suggestion}GiB'
|
||||
|
||||
return max_memory if len(max_memory) > 0 else None
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user