huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-11 20:34:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fixing config.n_head.

Browse Source
This commit is contained in:
Nicolas Patry 2024-07-02 17:01:25 +00:00
parent 24bbd7b822
commit e8ff76fd18
No known key found for this signature in database
GPG Key ID: E939E8CC91A1C674
1 changed files with 5 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -884,14 +884,17 @@ class FlashCausalLM(Model):
model = model_class(prefix, config, weights)
torch.distributed.barrier(group=self.process_group)
# VLM models define the config we care about in their text_config
text_config = getattr(config, "text_config", None)
if text_config is not None:
config = text_config
self.num_layers = config.num_hidden_layers
# Validation is done in the model itself
num_heads = getattr(config, "num_key_value_heads", config.n_head)
if num_kv_heads is None:
num_kv_heads = config.num_key_value_heads
num_kv_heads = getattr(config, "num_key_value_heads", None)
if num_kv_heads is None:
# Final overide for GPT2
num_kv_heads = config.n_head
self.num_kv_heads = num_kv_heads // self.process_group.size()
self.head_size = config.hidden_size // config.num_attention_heads