fix: cast q,k,v a dtype originale dopo RoPE
Browse files- modeling_quark.py +4 -2
modeling_quark.py
CHANGED
|
@@ -70,14 +70,16 @@ class GroupedQueryAttention(nn.Module):
|
|
| 70 |
|
| 71 |
def forward(self, x, attention_mask=None, **kwargs):
|
| 72 |
B, T, _ = x.shape
|
|
|
|
| 73 |
q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
|
| 74 |
k = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
|
| 75 |
v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
|
| 76 |
-
q, k = self.rope(q, k)
|
| 77 |
if self.n_groups > 1:
|
| 78 |
k = k.repeat_interleave(self.n_groups, dim=1)
|
| 79 |
v = v.repeat_interleave(self.n_groups, dim=1)
|
| 80 |
-
|
|
|
|
| 81 |
out = F.scaled_dot_product_attention(
|
| 82 |
q, k, v, attn_mask=None,
|
| 83 |
dropout_p=self.drop if self.training else 0.0,
|
|
|
|
| 70 |
|
| 71 |
def forward(self, x, attention_mask=None, **kwargs):
|
| 72 |
B, T, _ = x.shape
|
| 73 |
+
dtype = x.dtype # salva dtype originale (bfloat16)
|
| 74 |
q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
|
| 75 |
k = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
|
| 76 |
v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
|
| 77 |
+
q, k = self.rope(q, k) # RoPE porta q,k a float32
|
| 78 |
if self.n_groups > 1:
|
| 79 |
k = k.repeat_interleave(self.n_groups, dim=1)
|
| 80 |
v = v.repeat_interleave(self.n_groups, dim=1)
|
| 81 |
+
# Porta tutti allo stesso dtype per SDPA
|
| 82 |
+
q, k, v = q.to(dtype), k.to(dtype), v.to(dtype)
|
| 83 |
out = F.scaled_dot_product_attention(
|
| 84 |
q, k, v, attn_mask=None,
|
| 85 |
dropout_p=self.drop if self.training else 0.0,
|