fix: dtype cast in SDPA + tie_weights() esplicito

modeling_quark.py

@@ -46,8 +46,8 @@ class RotaryEmbedding(nn.Module):
         T = q.size(2)
         if T > self._max:
             self._build_cache(T)
-        cos = self.cos_cache[:, :, :T, :]
-        sin = self.sin_cache[:, :, :T, :]
+        cos = self.cos_cache[:, :, :T, :].to(dtype=q.dtype)
+        sin = self.sin_cache[:, :, :T, :].to(dtype=q.dtype)
         return q * cos + self._rotate_half(q) * sin,                k * cos + self._rotate_half(k) * sin
 
 
@@ -75,6 +75,10 @@ class GroupedQueryAttention(nn.Module):
         if self.n_groups > 1:
             k = k.repeat_interleave(self.n_groups, dim=1)
             v = v.repeat_interleave(self.n_groups, dim=1)
+        # Forza dtype uniforme prima di SDPA
+        dtype = q.dtype
+        k = k.to(dtype)
+        v = v.to(dtype)
         out = F.scaled_dot_product_attention(
             q, k, v, attn_mask=None,
             dropout_p=self.drop if self.training else 0.0,
@@ -110,7 +114,7 @@ class TransformerBlock(nn.Module):
 
 
 class QuarkPreTrainedModel(PreTrainedModel):
-    config_class     = QuarkConfig
+    config_class      = QuarkConfig
     base_model_prefix = "model"
     supports_gradient_checkpointing = False
 
@@ -153,6 +157,11 @@ class QuarkForCausalLM(QuarkPreTrainedModel):
     def set_output_embeddings(self, value):
         self.lm_head = value
 
+    def tie_weights(self):
+        """HF chiama questo metodo dopo il caricamento dei pesi."""
+        if self.config.tie_word_embeddings:
+            self.lm_head.weight = self.embed_tokens.weight
+
     def forward(
         self,
         input_ids      = None,
@@ -160,9 +169,9 @@ class QuarkForCausalLM(QuarkPreTrainedModel):
         labels         = None,
         **kwargs,
     ):
-        x      = self.embed_tokens(input_ids)
+        x = self.embed_tokens(input_ids)
         for layer in self.layers:
-            x  = layer(x, attention_mask=attention_mask)
+            x = layer(x, attention_mask=attention_mask)
         x      = self.norm(x)
         logits = self.lm_head(x)
 
@@ -179,7 +188,6 @@ class QuarkForCausalLM(QuarkPreTrainedModel):
     @torch.no_grad()
     def generate_text(self, input_ids, max_new_tokens=200, temperature=0.7,
                       top_p=0.9, eos_token_id=None):
-        """Generazione semplice senza dipendere da .generate() HF."""
         ctx = input_ids.clone()
         for _ in range(max_new_tokens):
             out    = self(ctx[:, -self.config.max_seq_len:])
@@ -187,7 +195,6 @@ class QuarkForCausalLM(QuarkPreTrainedModel):
             if temperature > 0:
                 logits /= temperature
                 probs   = F.softmax(logits, dim=-1)
-                # top-p
                 sorted_p, sorted_i = torch.sort(probs, descending=True)
                 cum_p = torch.cumsum(sorted_p, dim=-1)
                 remove = cum_p - sorted_p > top_p