Upload model

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1daa5083d32e5c14f1f024b693d30a7414cce7cf7814b9ce1dc1b0eb03dfb46b
-size 385823368
+oid sha256:1946078d3452868c8afe7c88aff1e9364a13ecd67361113eb7515f652a70c302
+size 409219512

modeling_t2.py

@@ -322,11 +322,7 @@ class TransformerPreTrainedModel(PreTrainedModel):
 class TransformerModel(TransformerPreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
-        # self.wte = nn.Embedding(config.vocab_size, config.hidden_size)
-        self.wte = nn.Sequential(
-            nn.Embedding(config.vocab_size, config.stage_0_hidden_size),
-            StageLinear(config.stage_0_hidden_size, config.hidden_size, bias=False, stage=0, config=config)
-        )
+        self.wte = nn.Embedding(config.vocab_size, config.hidden_size)
         self.h = nn.ModuleList(
             [TransformerBlock(config, stage=1) for i in range(config.num_hidden_layers)]
         )
@@ -337,10 +333,10 @@ class TransformerModel(TransformerPreTrainedModel):
         self.post_init()
 
     def get_input_embeddings(self):
-        return self.wte[0]
+        return self.wte
 
     def set_input_embeddings(self, new_embeddings):
-        self.wte[0] = new_embeddings
+        self.wte = new_embeddings
 
     def forward(
         self,
@@ -500,7 +496,7 @@ class TransformerModel(TransformerPreTrainedModel):
         )
 
 class TransformerModelForCausalLM(TransformerPreTrainedModel):
-    _tied_weights_keys = ["lm_head.1.weight"]
+    _tied_weights_keys = ["lm_head.weight"]
     _tied_weights_keys = []
     def __init__(self, config):
         super().__init__(config)
@@ -508,20 +504,16 @@ class TransformerModelForCausalLM(TransformerPreTrainedModel):
         # self.lm_head = nn.Linear(
         #     config.hidden_size, config.vocab_size, bias=False
         # )
-        self.lm_head = nn.Sequential(
-            StageLinear(config.hidden_size, config.stage_0_hidden_size, bias=False, stage=0, config=config),
-            nn.Linear(config.stage_0_hidden_size, config.vocab_size),
-        )
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size)
         self.model_parallel = False
         self.device_map = None
         self.post_init()
 
     def get_output_embeddings(self):
-        return self.lm_head[1]
+        return self.lm_head
 
     def set_output_embeddings(self, new_embeddings):
-        # print("Huggingface is inexplicably trying to tie the lm head no matter how many times i'm saying tie_weights False but I'm not letting it")
-        self.lm_head[1] = new_embeddings
+        self.lm_head = new_embeddings
 
     def prepare_inputs_for_generation(
         self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs