From 88b0b6f4f13461bcf6f35a16c8d290cdf90c329c Mon Sep 17 00:00:00 2001 From: kingbri Date: Sun, 21 Apr 2024 23:47:37 -0400 Subject: [PATCH] Model: Cast autosplit_reserve to int Torch errors if float values are passed (because bytes are not float types). Therefore, overestimate and cast to an int type. Resolves #97 Signed-off-by: kingbri --- backends/exllamav2/model.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/backends/exllamav2/model.py b/backends/exllamav2/model.py index 36754f8..3cf4400 100644 --- a/backends/exllamav2/model.py +++ b/backends/exllamav2/model.py @@ -1,6 +1,7 @@ """The model container class for ExLlamaV2 models.""" import gc +import math import pathlib import threading import time @@ -130,7 +131,10 @@ class ExllamaV2Container: autosplit_reserve_megabytes = unwrap(kwargs.get("autosplit_reserve"), [96]) self.autosplit_reserve = list( - map(lambda value: value * 1024**2, autosplit_reserve_megabytes) + map( + lambda value: int(math.ceil(value * 1024**2)), + autosplit_reserve_megabytes, + ) ) elif gpu_count > 1: # Manual GPU split