Backends: Read autosplit_reserve for every exllamav3 split mode

backends/exllamav3/model.py

@@ -260,6 +260,15 @@ async def create(cls, model_directory: pathlib.Path, hf_model: HFModel, **kwargs
         gpu_device_list = list(range(0, gpu_count))
         use_tp = unwrap(kwargs.get("tensor_parallel"), False)
 
+        # Reserve VRAM per GPU. Read for every split mode
+        default_reserve = [96]
+        autosplit_reserve_megabytes = unwrap(kwargs.get("autosplit_reserve"), default_reserve)
+        if isinstance(autosplit_reserve_megabytes, (int, float)) and not isinstance(
+            autosplit_reserve_megabytes, bool
+        ):
+            autosplit_reserve_megabytes = [autosplit_reserve_megabytes]
+        self.autosplit_reserve = [value / 1024 for value in autosplit_reserve_megabytes]
+
         # Set GPU split options
         if gpu_count == 1:
             self.gpu_split_auto = False
@@ -286,21 +295,22 @@ async def create(cls, model_directory: pathlib.Path, hf_model: HFModel, **kwargs
                 self.gpu_split_auto = False
                 self.gpu_split = gpu_split
 
-                # Causes crash if set with GPU split
-                # TODO: Remove when fixed in exllama upstream
-                self.autosplit_reserve = None
-
                 gpu_device_list = [
                     device_idx for device_idx, memory in enumerate(self.gpu_split) if memory > 0
                 ]
             elif gpu_split_auto and not self.use_tp:
                 # Otherwise fallback to autosplit settings
                 self.gpu_split_auto = gpu_split_auto
 
-                autosplit_reserve_megabytes = unwrap(kwargs.get("autosplit_reserve"), [96])
-
-                # Reserve VRAM for each GPU
-                self.autosplit_reserve = [value / 1024 for value in autosplit_reserve_megabytes]
+        # A manual split states how much of each GPU to use, and exllamav3 takes
+        # either a reserve or a split, so the reserve is dropped for that model
+        if autosplit_reserve_megabytes != default_reserve:
+            if self.gpu_split:
+                xlogger.warning("autosplit_reserve is ignored when gpu_split is set.")
+            if self.draft_gpu_split:
+                xlogger.warning(
+                    "autosplit_reserve is ignored for the draft model when draft_gpu_split is set."
+                )
 
         if not hardware_supports_exllamav3(gpu_device_list):
             gpu_unsupported_message = (
@@ -686,9 +696,11 @@ async def load_gen(self, progress_callback=None, **kwargs):
 
     @torch.inference_mode()
     def load_model_sync(self, progress_callback=None):
+        # exllamav3 asserts on reserve_per_device and use_per_device together,
+        # so a model with a manual split gets the split and not the reserve
         if self.use_vision:
             for value in self.vision_model.load_gen(
-                reserve_per_device=self.autosplit_reserve,
+                reserve_per_device=None if self.gpu_split else self.autosplit_reserve,
                 use_per_device=self.gpu_split or None,
                 callback=progress_callback,
             ):
@@ -697,7 +709,9 @@ def load_model_sync(self, progress_callback=None):
 
         if self.use_draft_model:
             for value in self.draft_model.load_gen(
-                reserve_per_device=self.autosplit_reserve,
+                reserve_per_device=(
+                    None if (self.gpu_split or self.draft_gpu_split) else self.autosplit_reserve
+                ),
                 use_per_device=self.draft_gpu_split or None,
                 callback=progress_callback,
             ):
@@ -716,7 +730,7 @@ def load_model_sync(self, progress_callback=None):
         for value in self.model.load_gen(
             tensor_p=self.use_tp,
             tp_backend=self.tp_backend,
-            reserve_per_device=self.autosplit_reserve,
+            reserve_per_device=None if self.gpu_split else self.autosplit_reserve,
             use_per_device=self.gpu_split,
             callback=progress_callback,
             max_chunk_size=self.chunk_size,

common/config_models.py

@@ -236,8 +236,12 @@ class ModelConfig(BaseConfigModel):
     autosplit_reserve: List[float] = Field(
         [96],
         description=(
-            "Reserve VRAM used for autosplit loading (default: 96 MB on GPU 0).\n"
-            "Represented as an array of MB per GPU."
+            "Reserve VRAM used when loading a model (default: 96 MB on GPU 0).\n"
+            "Represented as an array of MB per GPU.\n"
+            "A negative value excludes that GPU from the model split, so\n"
+            "excluding every GPU will fail to load.\n"
+            "Ignored for a model whose placement is already set by gpu_split\n"
+            "or draft_gpu_split."
         ),
     )
     gpu_split: List[float] = Field(

config_sample.yml

@@ -119,8 +119,12 @@ model:
   # Not parsed for single GPU users.
   gpu_split_auto: true
 
-  # Reserve VRAM used for autosplit loading (default: 96 MB on GPU 0).
+  # Reserve VRAM used when loading a model (default: 96 MB on GPU 0).
   # Represented as an array of MB per GPU.
+  # A negative value excludes that GPU from the model split, so
+  # excluding every GPU will fail to load.
+  # Ignored for a model whose placement is already set by gpu_split
+  # or draft_gpu_split.
   autosplit_reserve: [96]
 
   # Array of VRAM sizes to split between GPUs, in GB (default: []).