Fix num_key_value_groups

9 months ago · cda4fe8cef
parent cac654a16f
commit cda4fe8cef
1 changed files with 3 additions and 1 deletions
--- a/src/petals/models/falcon/config.py
+++ b/src/petals/models/falcon/config.py
@ -21,8 +21,10 @@ class DistributedFalconConfig(DefaultRevisionMixin, FalconConfig, ClientConfig,

    @property
    def num_key_value_groups(self) -> int:
-        if self.new_decoder_architecture or not self.multi_query:
+        if self.new_decoder_architecture:
            return self.num_attention_heads // self.num_kv_heads
+        if self.multi_query:
+            return self.num_attention_heads
        return 1

    @classmethod