From dee649c60ce0afd6ee4cb1718b6882e75e1ba9d8 Mon Sep 17 00:00:00 2001 From: Nicolas Patry Date: Tue, 23 Jul 2024 11:26:53 +0200 Subject: [PATCH] Chore: Fix naming issues regarding head_size, there can only be one. --- .../models/custom_modeling/flash_mistral_modeling.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/server/text_generation_server/models/custom_modeling/flash_mistral_modeling.py b/server/text_generation_server/models/custom_modeling/flash_mistral_modeling.py index eeb3c45f..eca01bbb 100644 --- a/server/text_generation_server/models/custom_modeling/flash_mistral_modeling.py +++ b/server/text_generation_server/models/custom_modeling/flash_mistral_modeling.py @@ -149,15 +149,14 @@ class MistralAttention(torch.nn.Module): bias=False, ) - head_size = config.hidden_size // config.num_attention_heads self.query_key_value = TensorParallelMultiAdapterLinear.load( query_key_value, layer_id, ["q_proj", "k_proj", "v_proj"], sizes=[ - head_size * config.num_attention_heads, - head_size * config.num_key_value_heads, - head_size * config.num_key_value_heads, + self.head_size * config.num_attention_heads, + self.head_size * config.num_key_value_heads, + self.head_size * config.num_key_value_heads, ], process_group=weights.process_group, )