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

Merge pull request #18 from DocShotgun/main

Browse Source 
Add automatic NTK-aware alpha scaling to model
This commit is contained in:
Brian Dashore
2023-12-03 01:06:50 -05:00
committed by GitHub
parent 61f6e51fdb 27fc0c0069
commit 109e4223e0
1 changed files with 16 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
model.py
View File

@@ -70,11 +70,17 @@ class ModelContainer:
self.config.model_dir = str(model_directory.resolve())
self.config.prepare()
# Grab the base model's sequence length before overrides for rope calculations
base_seq_len = self.config.max_seq_len
# Then override the max_seq_len if present
if "max_seq_len" in kwargs: self.config.max_seq_len = kwargs["max_seq_len"]
if "rope_scale" in kwargs: self.config.scale_pos_emb = kwargs["rope_scale"]
if "rope_alpha" in kwargs: self.config.scale_alpha_value = kwargs["rope_alpha"]
if "no_flash_attn" in kwargs: self.config.no_flash_attn = kwargs["no_flash_attn"]
# Automatically calculate rope alpha
self.config.scale_alpha_value = kwargs.get("rope_alpha") or self.calculate_rope_alpha(base_seq_len)
if "no_flash_attn" in kwargs: self.config.no_flash_attn = kwargs["no_flash_attn"]
if "low_mem" in kwargs and kwargs["low_mem"]:
self.config.set_low_mem()
@@ -101,13 +107,7 @@ class ModelContainer:
self.draft_config.model_dir = str(draft_model_path.resolve())
self.draft_config.prepare()
if "draft_rope_alpha" in kwargs:
self.draft_config.scale_alpha_value = kwargs.get("draft_rope_alpha") or 1
else:
ratio = self.config.max_seq_len / self.draft_config.max_seq_len
alpha = -0.13436 + 0.80541 * ratio + 0.28833 * ratio ** 2
self.draft_config.scale_alpha_value = alpha
self.draft_config.scale_alpha_value = kwargs.get("draft_rope_alpha") or self.calculate_rope_alpha(self.draft_config.max_seq_len)
self.draft_config.max_seq_len = self.config.max_seq_len
if "chunk_size" in kwargs:
@@ -115,6 +115,13 @@ class ModelContainer:
self.draft_config.max_attn_size = kwargs["chunk_size"] ** 2
def calculate_rope_alpha(self, base_seq_len):
ratio = self.config.max_seq_len / base_seq_len
# Default to a 1 alpha if the sequence length is ever less than or equal to 1
alpha = 1 if ratio <= 1.0 else -0.13436 + 0.80541 * ratio + 0.28833 * ratio ** 2
return alpha
def get_model_path(self):
model_path = pathlib.Path(self.config.model_dir)
return model_path