Initial commit of GPT-2 46M SwiGLU Dense baseline model architecture and configs

configuration_gpt2.py

@@ -1,7 +1,11 @@
 from transformers import PretrainedConfig
 
-class GPT2Config(PretrainedConfig):
-    model_type = "gpt2"
+class GPT2CustomConfig(PretrainedConfig):
+    model_type = "gpt2_custom"
+    auto_map = {
+        "AutoConfig": "configuration_gpt2.GPT2CustomConfig",
+        "AutoModelForCausalLM": "modeling_gpt2.GPT2CustomLMHeadModel"
+    }
     
     def __init__(
         self,

modeling_gpt2.py

@@ -4,7 +4,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from transformers import PreTrainedModel, GenerationMixin, AutoConfig, AutoModelForCausalLM
-from configuration_gpt2 import GPT2Config
+from configuration_gpt2 import GPT2CustomConfig
 
 class CausalSelfAttention(nn.Module):
     def __init__(self, config):
@@ -62,7 +62,7 @@ class GPT2Block(nn.Module):
         return x
 
 class GPT2CustomLMHeadModel(PreTrainedModel, GenerationMixin):
-    config_class = GPT2Config
+    config_class = GPT2CustomConfig
     base_model_prefix = "transformer"
     
     def __init__(self, config):
@@ -114,5 +114,5 @@ class GPT2CustomLMHeadModel(PreTrainedModel, GenerationMixin):
         return {"input_ids": input_ids}
 
 # Register configuration and model for auto mapping
-AutoConfig.register("gpt2", GPT2Config)
-AutoModelForCausalLM.register(GPT2Config, GPT2CustomLMHeadModel)
+AutoConfig.register("gpt2_custom", GPT2CustomConfig)
+AutoModelForCausalLM.register(GPT2CustomConfig, GPT2CustomLMHeadModel)

train.py

@@ -75,7 +75,7 @@ def run_pipeline(model_name: str, epochs: int = 10, skip_eval: bool = False, ski
     
     # Import model architecture
     from modeling_gpt2 import GPT2CustomLMHeadModel
-    from configuration_gpt2 import GPT2Config
+    from configuration_gpt2 import GPT2CustomConfig
 
     # Print GPU details
     if torch.cuda.is_available():
@@ -232,7 +232,7 @@ def run_pipeline(model_name: str, epochs: int = 10, skip_eval: bool = False, ski
     steps_per_epoch  = chunks_per_epoch // GRAD_ACCUM_STEPS 
     total_steps      = steps_per_epoch * EPOCHS
 
-    cfg = GPT2Config(
+    cfg = GPT2CustomConfig(
         vocab_size=VOCAB_SIZE,
         n_positions=BLOCK_SIZE,
         n_embd=768,