Falcon/DBRX: get correct number of key-value heads (#2205)

server/text_generation_server/models/__init__.py

@@ -797,6 +797,10 @@ def get_model(
                     quantize=quantize,
                     speculator=speculator,
                     dtype=dtype,
+                    aliases={
+                        "lm_head.weight": ["transformer.word_embeddings.weight"],
+                        "transformer.word_embeddings.weight": ["lm_head.weight"],
+                    },
                     trust_remote_code=trust_remote_code,
                     lora_adapter_ids=lora_adapter_ids,
                     config_class=RWConfig,

server/text_generation_server/models/custom_modeling/flash_dbrx_modeling.py

@@ -105,6 +105,12 @@ class DbrxFFNConfig(PretrainedConfig):
 
 
 class DbrxConfig(PretrainedConfig):
+    attribute_map = {
+        "hidden_size": "d_model",
+        "num_attention_heads": "n_heads",
+        "num_hidden_layers": "n_layers",
+    }
+
     def __init__(
         self,
         d_model: int = 2048,
@@ -157,6 +163,12 @@ class DbrxConfig(PretrainedConfig):
             **kwargs,
         )
 
+    @property
+    def num_key_value_heads(self):
+        # We can't use the attribute map, since this the number of KV
+        # heads is not top-level.
+        return self.attn_config.kv_n_heads
+
 
 def promote_scalar(x: torch.Tensor) -> torch.Tensor:
     return x.view(1) if len(x.size()) == 0 else x

server/text_generation_server/models/custom_modeling/flash_rw_modeling.py

@@ -42,6 +42,7 @@ class RWConfig(PretrainedConfig):
     attribute_map = {
         "num_hidden_layers": "n_layer",
         "num_attention_heads": "n_head",
+        "num_key_value_heads": "n_head_kv",
     }
 
     def __init__(

server/text_generation_server/models/flash_causal_lm.py

@@ -905,11 +905,10 @@ class FlashCausalLM(Model):
         self.num_layers = config.num_hidden_layers
         # Validation is done in the model itself
         if num_kv_heads is None:
-            # Order is important here.
+            num_kv_heads = getattr(config, "num_key_value_heads", None)
-            for attr in ["num_key_value_heads", "num_attention_heads", "n_head"]:
+            # GPT-2 workaround
-                num_kv_heads = getattr(config, attr, None)
+            if num_kv_heads is None:
-                if num_kv_heads is not None:
+                num_kv_heads = getattr(config, "n_head", None)
-                    break
         if num_kv_heads is None:
             raise ValueError("Cannot get the number of key/value heads")
         self.num_kv_heads = (