Update model.py

model.py

@@ -141,7 +141,7 @@ class FalconConfig(PretrainedConfig):
 
 def built_bloom_alibi(attention_mask, num_attention_heads):
     b, s = attention_mask.shape
-    cp2 = 2 ** math.ceil(math.log2(num_attention_heads))
+    cp2 = 2 ** math.floor(math.log2(num_attention_heads))
     base = jnp.asarray(
         2 ** (- (2 ** -(math.log2(cp2) - 3))), dtype=jnp.float32
     )
@@ -155,8 +155,8 @@ def built_bloom_alibi(attention_mask, num_attention_heads):
         extra_power = jnp.arange(1, 1 + 2 * num_rem_heads, 2, dtype=jnp.dtype)
         slops = jnp.concatenate([slops, jnp.power(extra_base, extra_power)], axis=0)
     arange_tensor = (((jnp.cumsum(attention_mask, axis=-1)) - 1) * attention_mask)[:, jnp.newaxis, :]
-    alibi = slops[..., jnp.newaxis] * arange_tensor
-    return alibi.reshape(b, num_attention_heads, 1, s)
+    alibi = slops[..., jnp.newaxis].astype(jnp.bfloat16) * arange_tensor
+    return alibi.reshape(b , num_attention_heads, 1, s)
 
 
 def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0,
@@ -225,6 +225,7 @@ class FlaxFalconAttention(nn.Module):
                  attention_mask: jnp.DeviceArray = None,
                  ):
         b, s, d = hidden_states.shape
+
         qkv = self.w_qkv(hidden_states)
         if not self.config.multi_query:
             q, k, v = jnp.split(qkv, 3, -1)