Fix modeling_gpt1.py

Few fixes:
- correctly norm before layers, as in the original Transformer paper
- fix model output
- fix LayerNorm instead of RMSNorm

modeling_gpt1.py

@@ -9,33 +9,15 @@ from transformers.modeling_outputs import (
     BaseModelOutput,
     CausalLMOutput,
 )
-from transformers.activations import ACT2FN
+from transformers.activations import get_activation
 
 from configuration_gpt1 import GPT1Config
 
 
-class GPT1RMSNorm(nn.Module):
-    def __init__(self, config: GPT1Config):
-        super().__init__()
-        self.config = config
-        self.weight = nn.Parameter(torch.ones(config.hidden_size))
-
-    def _norm(self, x):
-        std = x.pow(2).mean(dim=-1, keepdim=True).sqrt()
-        return x / (std + self.config.layer_norm_eps)
-
-    def forward(self, hidden_state):
-        input_dtype = hidden_state.dtype
-        # compute in float32, not in fp16, since normalization needs to be accurate
-        hidden_state = hidden_state.float()
-        output = self._norm(hidden_state).type_as(input_dtype)
-        return output * self.weight
-
-
 class GPT1MLP(nn.Module):
     def __init__(self, config: GPT1Config):
         super().__init__()
-        self.activation_fn = ACT2FN(config.hidden_act)
+        self.activation_fn = get_activation(config.hidden_act)
         self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
         self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
 
@@ -109,25 +91,27 @@ class GPT1DecoderLayer(nn.Module):
         self.attention = GPT1Attention(config)
         self.mlp = GPT1MLP(config)
 
-        self.attention_norm = GPT1RMSNorm(config)
-        self.mlp_norm = GPT1RMSNorm(config)
+        self.attention_norm = nn.LayerNorm(normalized_shape=config.hidden_size,
+                                           eps=config.layer_norm_eps)
+        self.mlp_norm = nn.LayerNorm(normalized_shape=config.hidden_size,
+                                     eps=config.layer_norm_eps)
 
         self.res_dropout = nn.Dropout(p=config.resid_pdrop)
 
     def forward(self, hidden_state, attn_mask):
         # attention
         residual = hidden_state
-        hidden_state = self.attention_norm(hidden_state)
         hidden_state = self.attention(hidden_state, attn_mask)
         hidden_state = self.res_dropout(hidden_state)
         hidden_state = residual + hidden_state
+        hidden_state = self.attention_norm(hidden_state)
 
         # feed forward fully connected
         residual = hidden_state
-        hidden_state = self.mlp_norm(hidden_state)
         hidden_state = self.mlp(hidden_state)
         hidden_state = self.res_dropout(hidden_state)
         hidden_state = residual + hidden_state
+        hidden_state = self.mlp_norm(hidden_state)
 
         return hidden_state
 
@@ -165,8 +149,6 @@ class GPT1Model(GPT1PreTrainedModel):
             [GPT1DecoderLayer(config) for _ in range(config.num_hidden_layers)]
         )
 
-        self.norm = GPT1RMSNorm(config)
-
         causal_mask = torch.full((1, config.max_position_embeddings, config.max_position_embeddings),
                                  fill_value=float('-inf'))
         self.register_buffer('causal_mask',
@@ -182,9 +164,9 @@ class GPT1Model(GPT1PreTrainedModel):
         self.embs = value
 
     def forward(self, input_ids, *args, **kwargs):
-        position_ids = torch.arange(input_ids.size()[-1],
+        position_ids = torch.arange(input_ids.size(-1),
                                     dtype=torch.long,
-                                    device=input_ids.device)
+                                    device=input_ids.device).unsqueeze_(0)
 
         input_embeds = self.embs(input_ids) # (bs, seq_len, dim)
         position_embeds = self.pos_emb(position_ids)
@@ -192,11 +174,10 @@ class GPT1Model(GPT1PreTrainedModel):
 
         causal_mask = self.causal_mask.to(dtype=input_embeds.dtype,
                                           device=input_embeds.device)
+
         for layer in self.layers:
             hidden_state = layer(hidden_state, attn_mask=causal_mask)
 
-        hidden_state = self.norm(hidden_state)
-
         return BaseModelOutput(
             last_hidden_state=hidden_state
         )
@@ -236,7 +217,8 @@ class GPT1ForCausalLM(GPT1PreTrainedModel):
     def forward(self, input_ids, labels = None, *args, **kwargs):
         output = self.model(input_ids)
 
-        logits = self.lm_head(output).float()
+        hidden_state = output[0]
+        logits = self.lm_head(hidden_state).float()
 
         loss = None
         if labels is not None: