theroyallab/tabbyAPI
1
0
Fork 0
mirror of https://github.com/theroyallab/tabbyAPI.git synced 2026-03-15 00:07:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #254 from lucyknada/main

Browse Source 
add draft_gpu_split option for spec decoding
This commit is contained in:
Brian
2025-02-11 16:48:03 -05:00
committed by GitHub
parent e290b88568 beb6d8faa5
commit 2e491472d1
3 changed files with 59 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

62
backends/exllamav2/model.py
View File

@@ -89,7 +89,8 @@ class ExllamaV2Container:
generation_config: Optional[GenerationConfig] = None
# GPU split vars
gpu_split: Optional[list] = None
gpu_split: List[float] = []
draft_gpu_split: List[float] = []
gpu_split_auto: bool = True
autosplit_reserve: List[float] = [96 * 1024**2]
use_tp: bool = False
@@ -180,6 +181,7 @@ class ExllamaV2Container:
)
draft_model_path = draft_model_path / draft_model_name
self.draft_gpu_split = draft_args.get("draft_gpu_split")
self.draft_model_dir = draft_model_path
self.draft_config.model_dir = str(draft_model_path.resolve())
self.draft_config.prepare()
@@ -232,6 +234,15 @@ class ExllamaV2Container:
for value in autosplit_reserve_megabytes
]
# Change the GPU device list only if gpu_split's list is too small
# This allows for an uneven list specification
if self.draft_gpu_split and len(self.draft_gpu_split) > len(self.gpu_split):
gpu_device_list = [
device_idx
for device_idx, memory in enumerate(self.draft_gpu_split)
if memory > 0
]
# Hardcode max output length to 16
self.config.max_output_len = 16
@@ -375,6 +386,7 @@ class ExllamaV2Container:
# Set draft cache mode
self.draft_cache_mode = unwrap(draft_args.get("draft_cache_mode"), "FP16")
# Edit the draft config size
if chunk_size:
self.draft_config.max_input_len = chunk_size
self.draft_config.max_attention_size = chunk_size**2
@@ -619,21 +631,41 @@ class ExllamaV2Container:
# Draft uses the autosplit loader, so create a cache that reflects this
draft_cache_class = self.get_cache_class(self.draft_cache_mode)
self.draft_cache = self.create_cache(
cache_class=draft_cache_class,
autosplit=True,
use_tp=False,
model=self.draft_model,
)
for value in self.draft_model.load_autosplit_gen(
self.draft_cache,
reserve_vram=autosplit_reserve,
last_id_only=True,
callback_gen=progress_callback,
):
if value:
yield value
if self.draft_gpu_split:
logger.info("Loading with a manual GPU split (or a one GPU setup)")
for value in self.draft_model.load_gen(
self.draft_gpu_split,
callback_gen=progress_callback,
):
if value:
yield value
self.draft_cache = self.create_cache(
cache_class=draft_cache_class,
autosplit=False,
use_tp=False,
model=self.draft_model,
)
else:
logger.info("Loading with autosplit")
self.draft_cache = self.create_cache(
cache_class=draft_cache_class,
autosplit=True,
use_tp=False,
model=self.draft_model,
)
for value in self.draft_model.load_autosplit_gen(
self.draft_cache,
reserve_vram=autosplit_reserve,
last_id_only=True,
callback_gen=progress_callback,
):
if value:
yield value
# Test VRAM allocation with a full-length forward pass
input_ids = torch.zeros((1, self.config.max_input_len), dtype=torch.long)