Add support for Grouped Query Attention on Llama Model (#393)

Resolves #388
huggingface · Nov 15, 2023 · 4e4148c · 4e4148c
1 parent 35d61f5
commit 4e4148c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/models.js b/src/models.js
@@ -3035,9 +3035,9 @@ export class LlamaPreTrainedModel extends PreTrainedModel {
         // config doesn't contain pad_token_id, so we assume it is the eos_token_id
         this.config.pad_token_id = this.config.eos_token_id
 
-        this.num_heads = this.config.num_attention_heads
+        this.num_heads = this.config.num_key_value_heads ?? this.config.num_attention_heads
         this.num_layers = this.config.num_hidden_layers
-        this.dim_kv = this.config.hidden_size / this.num_heads;
+        this.dim_kv = this.config.hidden_size / this.config.num_attention_heads
     }
 }
 /**