fix: cast q,k a dtype di v dopo RoPE — identico a train.py

modeling_quark.py

@@ -1,5 +1,5 @@
 """
-Quark language model — copia esatta dell'architettura di training.
+Quark-72M — wrapper HuggingFace che usa l'architettura originale di training.
 """
 import math
 import torch
@@ -10,6 +10,8 @@ from transformers.modeling_outputs import CausalLMOutputWithPast
 from .configuration_quark import QuarkConfig
 
 
+# ── Architettura identica a train.py ─────────────────────────────────────────
+
 class RMSNorm(nn.Module):
     def __init__(self, dim, eps=1e-5):
         super().__init__()
@@ -24,16 +26,14 @@ class RMSNorm(nn.Module):
 class RotaryEmbedding(nn.Module):
     def __init__(self, head_dim, max_seq_len, theta=10_000.0):
         super().__init__()
-        assert head_dim % 2 == 0
         inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
         self.register_buffer("inv_freq", inv_freq, persistent=True)
         self._build_cache(max_seq_len)
 
     def _build_cache(self, seq_len):
-        device = self.inv_freq.device
-        t      = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)
-        freqs  = torch.outer(t, self.inv_freq)
-        emb    = torch.cat([freqs, freqs], dim=-1)
+        t     = torch.arange(seq_len, device=self.inv_freq.device).float()
+        freqs = torch.outer(t, self.inv_freq)
+        emb   = torch.cat([freqs, freqs], dim=-1)
         self.register_buffer("cos_cache", emb.cos()[None, None], persistent=False)
         self.register_buffer("sin_cache", emb.sin()[None, None], persistent=False)
         self._max = seq_len
@@ -49,6 +49,7 @@ class RotaryEmbedding(nn.Module):
             self._build_cache(T)
         cos = self.cos_cache[:, :, :T, :]
         sin = self.sin_cache[:, :, :T, :]
+        # Identico a train.py — nessun cast, broadcast naturale
         q = q * cos + self._rotate_half(q) * sin
         k = k * cos + self._rotate_half(k) * sin
         return q, k
@@ -57,7 +58,6 @@ class RotaryEmbedding(nn.Module):
 class GroupedQueryAttention(nn.Module):
     def __init__(self, cfg):
         super().__init__()
-        assert cfg.n_heads % cfg.n_kv_heads == 0
         self.n_heads    = cfg.n_heads
         self.n_kv_heads = cfg.n_kv_heads
         self.n_groups   = cfg.n_heads // cfg.n_kv_heads
@@ -69,22 +69,21 @@ class GroupedQueryAttention(nn.Module):
         self.rope   = RotaryEmbedding(cfg.head_dim, cfg.max_seq_len, cfg.rope_theta)
         self.drop   = cfg.dropout
 
-    def forward(self, x, attention_mask=None, **kwargs):
+    def forward(self, x, **kwargs):
         B, T, _ = x.shape
-        dtype = x.dtype   # salva dtype originale (bfloat16)
         q = self.q_proj(x).view(B, T, self.n_heads,    self.head_dim).transpose(1, 2)
         k = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
         v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
-        q, k = self.rope(q, k)   # RoPE porta q,k a float32
+        q, k = self.rope(q, k)
         if self.n_groups > 1:
             k = k.repeat_interleave(self.n_groups, dim=1)
             v = v.repeat_interleave(self.n_groups, dim=1)
-        # Porta tutti allo stesso dtype per SDPA
-        q, k, v = q.to(dtype), k.to(dtype), v.to(dtype)
+        # Cast uniforme prima di SDPA (q/k possono essere float32 dopo RoPE)
+        dtype = v.dtype
+        q, k = q.to(dtype), k.to(dtype)
         out = F.scaled_dot_product_attention(
-            q, k, v, attn_mask=None,
+            q, k, v, is_causal=True,
             dropout_p=self.drop if self.training else 0.0,
-            is_causal=True,
         )
         out = out.transpose(1, 2).contiguous().view(B, T, self.n_heads * self.head_dim)
         return self.o_proj(out)
@@ -115,20 +114,18 @@ class TransformerBlock(nn.Module):
         return x
 
 
+# ── HuggingFace wrapper ───────────────────────────────────────────────────────
+
 class QuarkPreTrainedModel(PreTrainedModel):
     config_class      = QuarkConfig
     base_model_prefix = "model"
     _keys_to_ignore_on_load_missing = ["lm_head.weight"]
-    supports_gradient_checkpointing = False
 
     def _init_weights(self, module):
-        std = 0.02
-        if isinstance(module, nn.Linear):
-            nn.init.normal_(module.weight, 0.0, std)
-            if module.bias is not None:
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            nn.init.normal_(module.weight, 0.0, 0.02)
+            if hasattr(module, "bias") and module.bias is not None:
                 nn.init.zeros_(module.bias)
-        elif isinstance(module, nn.Embedding):
-            nn.init.normal_(module.weight, 0.0, std)
 
 
 class QuarkForCausalLM(QuarkPreTrainedModel):
@@ -139,25 +136,26 @@ class QuarkForCausalLM(QuarkPreTrainedModel):
         self.embed_tokens = nn.Embedding(config.vocab_size, config.d_model)
         self.layers       = nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layers)])
         self.norm         = RMSNorm(config.d_model, config.rms_eps)
-        # lm_head usa embed_tokens.weight (weight tying) — non è un parametro separato
+        self.lm_head      = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        self.lm_head.weight = self.embed_tokens.weight
         self.post_init()
 
     def get_input_embeddings(self):  return self.embed_tokens
     def set_input_embeddings(self, v): self.embed_tokens = v
-    def get_output_embeddings(self): return None
-    def tie_weights(self, **kwargs): pass
+    def get_output_embeddings(self): return self.lm_head
+    def set_output_embeddings(self, v): self.lm_head = v
+    def tie_weights(self, **kwargs): self.lm_head.weight = self.embed_tokens.weight
 
     def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
         x = self.embed_tokens(input_ids)
         for layer in self.layers:
             x = layer(x)
         x      = self.norm(x)
-        # Weight tying: proiezione finale con la stessa matrice degli embedding
-        logits = x @ self.embed_tokens.weight.T
+        logits = self.lm_head(x)
         loss   = None
         if labels is not None:
             loss = F.cross_entropy(
-                logits[:, :-1].contiguous().view(-1, self.config.vocab_size),
+                logits[:, :-1].contiguous().view(-1, config.vocab_size),
                 labels[:, 1:].contiguous().view(-1),
                 ignore_index=-100,
             )
@@ -170,25 +168,17 @@ class QuarkForCausalLM(QuarkPreTrainedModel):
         for _ in range(max_new_tokens):
             out    = self(ctx[:, -self.config.max_seq_len:])
             logits = out.logits[0, -1, :].float()
-            if logits.isnan().any() or logits.isinf().any():
-                logits = torch.zeros_like(logits)
-                logits[2] = 1.0  # forza </s>
-            if temperature <= 0:
+            if temperature <= 0 or logits.isnan().any():
                 token = logits.argmax().view(1, 1)
             else:
                 logits -= logits.max()
                 logits /= temperature
                 probs   = F.softmax(logits, dim=-1)
                 sorted_p, sorted_i = torch.sort(probs, descending=True)
-                cum_p  = torch.cumsum(sorted_p, dim=-1)
-                mask   = (cum_p - sorted_p) > top_p
-                sorted_p[mask] = 0.0
-                total  = sorted_p.sum()
-                if total <= 0:
-                    token = sorted_i[0].view(1, 1)
-                else:
-                    sorted_p /= total
-                    token = sorted_i[torch.multinomial(sorted_p, 1)].view(1, 1)
+                cum_p = torch.cumsum(sorted_p, dim=-1)
+                sorted_p[(cum_p - sorted_p) > top_p] = 0.0
+                total = sorted_p.sum()
+                token = sorted_i[torch.multinomial(sorted_p / (total if total > 0 else 1), 1)].view(1, 1)
             ctx = torch.cat([ctx, token], dim=1)
             if eos_token_id is not None and token.item() == eos_token_id:
                 break