End of training

modeling_gptoss_mini.py

@@ -98,6 +98,9 @@ class Block(nn.Module):
 class GPTMiniForCausalLM(PreTrainedModel, GenerationMixin):
     config_class = GPTMiniConfig
 
+    _keys_to_ignore_on_save = []
+    _dynamic_tied_weights_keys = {"lm_head.weight", "embed.weight"}
+    
     def __init__(self, config: GPTMiniConfig):
         super().__init__(config)
         self.embed = nn.Embedding(config.vocab_size, config.hidden_size)
@@ -106,6 +109,7 @@ class GPTMiniForCausalLM(PreTrainedModel, GenerationMixin):
         self.ln_f = RMSNorm(config.hidden_size)
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
 
+        self.lm_head.weight = self.embed.weight
         self.post_init()
 
     def get_input_embeddings(self):