fix: weight tying via embed_tokens.weight.T, rimuove lm_head

modeling_quark.py

@@ -138,24 +138,21 @@ 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)
-        self.lm_head      = nn.Linear(config.d_model, config.vocab_size, bias=False)
-        self.lm_head.weight = self.embed_tokens.weight
+        # lm_head usa embed_tokens.weight (weight tying) — non è un parametro separato
         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 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 get_output_embeddings(self): return None
+    def tie_weights(self, **kwargs): pass
 
     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)
-        logits = self.lm_head(x)
+        # Weight tying: proiezione finale con la stessa matrice degli embedding
+        logits = x @ self.embed_tokens.weight.T
         loss   = None
         if labels is not None:
             loss = F.cross_entropy(
@@ -172,6 +169,9 @@ 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:
                 token = logits.argmax().view(1, 1)
             else: