feat: remove hidden_dim == head_dim * num_heads constraint

src/neuronx_distributed_inference/models/llama/modeling_llama.py

@@ -588,7 +588,7 @@ def __init__(self, config: InferenceConfig, tensor_model_parallel_group=None):
         self.hidden_size = config.hidden_size
         self.num_attention_heads = config.num_attention_heads
         self.num_key_value_heads = config.num_key_value_heads
-        self.head_dim = self.hidden_size // self.num_attention_heads
+        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
         self.max_position_embeddings = config.max_position_embeddings
         self.rope_theta = config.rope_theta
         self.padding_side = config.neuron_config.padding_side

src/neuronx_distributed_inference/modules/attention/attention_base.py

@@ -91,12 +91,6 @@ def __init__(self, tensor_model_parallel_group: Optional[ProcessGroup] = None):
         self.o_proj_layer_name = "o_proj"
 
     def init_gqa_properties(self):
-        if (self.head_dim * self.num_attention_heads) != self.hidden_size:
-            raise ValueError(
-                f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
-                f" and `num_heads`: {self.num_attention_heads})."
-            )
-
         self.qkv_proj = GroupQueryAttention_QKV(
             hidden_size=self.hidden_size,
             head_dim=self.head_dim,

src/neuronx_distributed_inference/modules/kvcache/kv_cache_manager.py

@@ -93,6 +93,9 @@ def _get_num_kv_heads_per_rank(self, config: InferenceConfig):
         return num_kv_heads_per_rank
 
     def _get_hidden_dim_per_head(self, config: InferenceConfig):
+        if hasattr(config, "head_dim"):
+            return config.head_dim
+
         hidden_size = config.hidden_size
         num_atten_head = config.num_attention_heads
         hidden_dim_per_head = hidden_size // num_atten_head