support stablm2 12b (#11265)

This commit is contained in:
Yishuo Wang 2024-06-07 15:46:00 +08:00 committed by GitHub
parent dbc3c2d72d
commit 42fab480ea
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -131,6 +131,10 @@ def stablelm_attention_forward(
query_states, key_states, value_states = qkv.split([self.num_heads, query_states, key_states, value_states = qkv.split([self.num_heads,
self.num_key_value_heads, self.num_key_value_heads,
self.num_key_value_heads], dim=1) self.num_key_value_heads], dim=1)
# For stablelm-2-12b's qk per-head norm
if getattr(self, "qk_layernorm", False):
query_states = self.q_layernorm(query_states)
key_states = self.k_layernorm(key_states)
kv_seq_len = key_states.shape[-2] kv_seq_len = key_states.shape[-2]
if past_key_value is not None: if past_key_value is not None: