allow custom head_dim for qwen2_moe (#37188)

allow custom head_dim

Co-authored-by: ryan.agile <ryan.agile@kakaobrain.com>
Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>
This commit is contained in:
Minho Ryu 2025-06-04 21:27:30 +09:00 committed by GitHub
parent 82fa68ca14
commit 6c5d4b1dd2
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -281,7 +281,7 @@ class Qwen2MoeAttention(nn.Module):
self.hidden_size = config.hidden_size
self.num_heads = config.num_attention_heads
self.head_dim = self.hidden_size // self.num_heads
self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
self.num_key_value_heads = config.num_key_value_heads
self.num_key_value_groups = self.num_heads // self.num_key_value_heads
self.max_position_embeddings = config.max_position_embeddings