Upload modeling_quark.py with huggingface_hub

modeling_quark.py

@@ -74,13 +74,22 @@ class QuarkForCausalLM(QuarkPreTrainedModel):
         self.layers=nn.ModuleList([QuarkBlock(config) for _ in range(config.n_layers)])
         self.norm=QuarkRMSNorm(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
+        self.lm_head.weight=self.embed_tokens.weight  # weight tying
         self.post_init()
-        self.lm_head.weight = self.model.embed_tokens.weight
+        
+    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
+        """Se lm_head.weight manca, copia da embed_tokens.weight (weight tying)"""
+        lm_key = f"{prefix}lm_head.weight"
+        emb_key = f"{prefix}embed_tokens.weight"
+        if lm_key not in state_dict and emb_key in state_dict:
+            state_dict[lm_key] = state_dict[emb_key].clone()
+        super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
+        
     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 forward(self, input_ids, attention_mask=None, labels=None, **kwargs):
         h=self.embed_tokens(input_ids)
         for layer in self.layers: h=layer(h)
@@ -90,4 +99,5 @@ class QuarkForCausalLM(QuarkPreTrainedModel):
             loss=F.cross_entropy(logits[...,:-1,:].contiguous().view(-1,self.config.vocab_size),
                                  labels[...,1:].contiguous().view(-1),ignore_index=-100)
         return CausalLMOutputWithPast(loss=loss, logits=logits)
+        
     def prepare_inputs_for_generation(self, input_ids, **kwargs): return {"input_ids": input_ids}