Fix modeling and configuration scripts

Replace modeling with an actual implementation, instead of using
PyTorch's built-in Transformer. This should be easier to configure and
hack.

configuration_gpt1.py

@@ -1,6 +1,6 @@
 """ GPT1 model configuration """
 
-from transformers import PretrainedConfig
+from transformers.configuration_utils import PretrainedConfig
 
 
 class GPT1Config(PretrainedConfig):

modeling_gpt1.py

@@ -1,10 +1,137 @@
+""" PyTorch GPT1 model."""
+
+import math
+
 import torch
 from torch import nn
 from transformers import PreTrainedModel
+from transformers.modeling_outputs import (
+    BaseModelOutput,
+    CausalLMOutput,
+)
+from transformers.activations import ACT2FN
 
 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.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
+        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
+
+    def forward(self, hidden_state):
+        hidden_state = self.fc1(hidden_state)
+        hidden_state = self.activation_fn(hidden_state)
+        hidden_state = self.fc2(hidden_state)
+        return hidden_state
+
+
+class GPT1Attention(nn.Module):
+    def __init__(self, config: GPT1Config):
+        """
+        Multi-head attention layer.
+        """
+        super().__init__()
+
+        assert config.hidden_size % config.num_attention_heads == 0
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        self.attn_dropout = nn.Dropout(p=config.attention_dropout)
+
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim)
+        self.k_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size)
+
+    def forward(self, hidden_state, attn_mask):
+        bs, seq_len, _ = hidden_state.size() # (batch_size, seq_len, dim)
+
+        # linearly project the inputs
+        Q = self.q_proj(hidden_state) # (batch_size, seq_len, n_heads * head_dim)
+        K = self.k_proj(hidden_state)
+        V = self.v_proj(hidden_state)
+
+        # split into n_heads to compute attention
+        queries = Q.view(bs, seq_len, self.num_heads, self.head_dim).transpose(1, 2) # (batch_size, n_heads, seq_len, head_dim)
+        keys = K.view(bs, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        values = V.view(bs, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+
+        # compute attention matmul
+        keys = keys.transpose(2, 3) # (batch_size, n_heads, head_dim, seq_len)
+        attn_scores = queries @ keys # (batch_size, n_heads, seq_len, seq_len)
+
+        # scale
+        attn_scores = attn_scores / math.sqrt(self.head_dim)
+
+        # mask
+        if attn_mask is not None:
+            attn_scores = attn_scores + attn_mask
+
+        # softmax (attention probabilities) + dropout
+        attn_probs = nn.functional.softmax(attn_scores, dim=-1, dtype=torch.float32).to(Q.dtype)
+        attn_probs = self.attn_dropout(attn_probs)
+
+        # matmul
+        attn_output = attn_probs @ values # (batch_size, n_heads, seq_len, head_dim)
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(bs, seq_len, self.hidden_size) # (batch_size, seq_len, n_heads * head_dim)
+
+        # final linear
+        attn_output = self.o_proj(attn_output)
+        return attn_output
+
+
+class GPT1DecoderLayer(nn.Module):
+    def __init__(self, config: GPT1Config):
+        super().__init__()
+        self.attention = GPT1Attention(config)
+        self.mlp = GPT1MLP(config)
+
+        self.attention_norm = GPT1RMSNorm(config)
+        self.mlp_norm = GPT1RMSNorm(config)
+
+        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
+
+        # 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
+
+        return hidden_state
+
+
 class GPT1PreTrainedModel(PreTrainedModel):
     config_class = GPT1Config
     supports_gradient_checkpointing = False
@@ -27,50 +154,55 @@ class GPT1Model(GPT1PreTrainedModel):
         super().__init__(config)
 
         # embeddings
-        self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size)
-        self.embeddings_dropout = nn.Dropout(config.embd_pdrop)
+        self.embs = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.embs_dropout = nn.Dropout(p=config.embd_pdrop)
 
         # positional encoding (learned)
         self.pos_emb = nn.Embedding(config.max_position_embeddings,
                                     config.hidden_size)
 
-        dec_layers = nn.TransformerEncoderLayer(d_model=config.hidden_size,
-                                                nhead=config.num_attention_heads,
-                                                dim_feedforward=config.intermediate_size,
-                                                dropout=config.attention_dropout,
-                                                activation=config.hidden_act,
-                                                layer_norm_eps=config.layer_norm_eps,
-                                                batch_first=True)
+        self.layers = nn.ModuleList(
+            [GPT1DecoderLayer(config) for _ in range(config.num_hidden_layers)]
+        )
+
+        self.norm = GPT1RMSNorm(config)
 
-        self.layers = nn.TransformerEncoder(dec_layers, config.num_hidden_layers)
+        causal_mask = torch.full((1, config.max_position_embeddings, config.max_position_embeddings),
+                                 fill_value=float('-inf'))
+        self.register_buffer('causal_mask',
+                             torch.triu(causal_mask, diagonal=1),
+                             persistent=False)
 
         self.post_init()
 
-    def forward(
-        self,
-        input_ids,
-        attention_mask
-    ):
+    def get_input_embeddings(self):
+        return self.embs
+
+    def set_input_embeddings(self, value):
+        self.embs = value
+
+    def forward(self, input_ids, *args, **kwargs):
         position_ids = torch.arange(input_ids.size()[-1],
                                     dtype=torch.long,
                                     device=input_ids.device)
 
-        input_embeds = self.embeddings(input_ids) # (bs, seq_len, dim)
+        input_embeds = self.embs(input_ids) # (bs, seq_len, dim)
         position_embeds = self.pos_emb(position_ids)
-        hidden_state = input_embeds + position_embeds
-
-        hidden_state = self.embeddings_dropout(hidden_state)
-
-        _, seq_len, _ = hidden_state.shape # (bs, seq_len, dim)
+        hidden_state = self.embs_dropout(input_embeds) + position_embeds
 
-        attention_mask = nn.Transformer.generate_square_subsequent_mask(seq_len).to(hidden_state.device)
+        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)
 
-        output = self.layers(hidden_state, attention_mask, is_causal=True)
+        hidden_state = self.norm(hidden_state)
 
-        return output
+        return BaseModelOutput(
+            last_hidden_state=hidden_state
+        )
 
 
-class GPT1ModelForCausalLM(GPT1PreTrainedModel):
+class GPT1ForCausalLM(GPT1PreTrainedModel):
     _tied_weights_keys = ["lm_head.weight"]
 
     def __init__(self, config: GPT1Config):
@@ -80,15 +212,29 @@ class GPT1ModelForCausalLM(GPT1PreTrainedModel):
 
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size)
 
+        # initialize weigths and apply final processing
         self.post_init()
 
-    def forward(
-        self,
-        input_ids,
-        attention_mask,
-        labels
-    ):
-        output = self.model(input_ids, attention_mask)
+    def get_input_embeddings(self):
+        return self.model.embs
+
+    def set_input_embeddings(self, value):
+        self.model.embs = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def get_decoder(self):
+        return self.model
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def forward(self, input_ids, labels = None, *args, **kwargs):
+        output = self.model(input_ids)
 
         logits = self.lm_head(output).float()
 
@@ -97,9 +243,15 @@ class GPT1ModelForCausalLM(GPT1PreTrainedModel):
             shift_logits = logits[..., :-1, :].contiguous()
             shift_labels = labels[..., 1:].contiguous()
 
-            loss_fct = torch.nn.CrossEntropyLoss()
+            loss_fn = torch.nn.CrossEntropyLoss()
             shift_logits = shift_logits.view(-1, self.config.vocab_size)
             shift_labels = shift_labels.view(-1)
-            loss = loss_fct(shift_logits, shift_labels)
-            return { "loss": loss, "logits": logits }
-        return { "logits": logits }
+            loss = loss_fn(shift_logits, shift_labels)
+
+        return CausalLMOutput(
+            loss=loss,
+            logits=logits
+        )
+
+    def prepare_inputs_for_generation(self, input_ids, *args, **kwargs):
+        return { 'input_ids': input_ids }