tiiuae

/

falcon-7b

like 1.05k

For some reason, the class "Attention" module produces a shape mismatch error when it goes to this branch as below. It complains that self.num_heads is not equal to self.num_kv, which will happen because the model config default value for num_kv is 1.

    if alibi is None:
        query_layer_ = query_layer.reshape(batch_size, self.num_heads, -1, self.head_dim)
        key_layer_ = key_layer.reshape(batch_size, self.num_kv, -1, self.head_dim)
        value_layer_ = value_layer.reshape(batch_size, self.num_kv, -1, self.head_dim).expand(-1, self.num_heads, -1, -1)
         attn_output = F.scaled_dot_product_attention(
            query_layer_, key_layer_, value_layer_, None, 0.0, is_causal=True
        )