huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-06-19 07:42:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

improve headdim

Browse Source
This commit is contained in:
Mohit Sharma 2025-04-24 09:55:14 +00:00
parent b86a73d72b
commit 36c5ec2abe
2 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
server/text_generation_server/models/custom_modeling/flash_mistral_modeling.py
View File

@ -116,7 +116,7 @@ class MistralAttention(torch.nn.Module):
) )
self.num_heads = config.num_attention_heads self.num_heads = config.num_attention_heads
self.hidden_size = config.hidden_size self.hidden_size = config.hidden_size
if hasattr(config, "head_dim") and config.head_dim is not None: if getattr(config, "head_dim", None) is not None:
self.head_size = config.head_dim self.head_size = config.head_dim
else: else:
self.head_size = self.hidden_size // self.num_heads self.head_size = self.hidden_size // self.num_heads

2
server/text_generation_server/models/flash_causal_lm.py
View File

@ -1298,7 +1298,7 @@ class FlashCausalLM(Model):
if head_size is None: if head_size is None:
# Some models use GQA and different sizes for o_proj # Some models use GQA and different sizes for o_proj
# and q_proj, that allows for that. # and q_proj, that allows for that.
if hasattr(config, "head_dim") and config.head_dim is not None: if getattr(config, "head_dim", None) is not None:
self.head_size = config.head_dim self.head_size = config.head_dim
else: else:
self.head_size = config.hidden_size // config.num_attention_heads self.head_size = config.hidden_size // config.num_attention_heads