Added files for quantization

aurora_model.py

@@ -0,0 +1,104 @@
+import torch
+import torch.nn as nn
+from transformers.models.gpt2.modeling_gpt2 import GPT2PreTrainedModel, GPT2Model
+from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions
+from transformers import PretrainedConfig
+
+
+class AuroraGPT2Config(PretrainedConfig):
+    model_type = "aurora-gpt2"
+
+    def __init__(
+        self,
+        vocab_size=25001,
+        n_positions=256,
+        n_embd=256,
+        n_layer=4,
+        n_head=4,
+        activation_function="gelu_new",
+        attn_pdrop=0.1,
+        embd_pdrop=0.1,
+        resid_pdrop=0.1,
+        layer_norm_epsilon=1e-05,
+        bos_token_id=50256,
+        eos_token_id=50256,
+        initializer_range=0.02,
+        use_cache=True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.n_positions = n_positions
+        self.max_position_embeddings = n_positions
+        self.n_embd = n_embd
+        self.hidden_size = n_embd   # <-- add this line (fix!)
+        self.n_layer = n_layer
+        self.num_hidden_layers = n_layer  # <-- add alias for HF
+        self.n_head = n_head
+        self.num_attention_heads = n_head  # <-- add alias for HF
+        self.activation_function = activation_function
+        self.attn_pdrop = attn_pdrop
+        self.embd_pdrop = embd_pdrop
+        self.resid_pdrop = resid_pdrop
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id
+        self.initializer_range = initializer_range
+        self.use_cache = use_cache
+
+class AuroraGPT2ForCausalLM(GPT2PreTrainedModel):
+    config_class = AuroraGPT2Config
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.transformer = GPT2Model(config)
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+
+        # Tie embeddings and output projection
+        self.transformer.wte.weight = self.lm_head.weight
+
+        # Initialize weights
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        labels=None,
+        **kwargs,
+    ):
+        # Run the transformer
+        outputs = self.transformer(
+            input_ids,
+            attention_mask=attention_mask,
+            **kwargs,
+        )
+
+        hidden_states = outputs.last_hidden_state
+        logits = self.lm_head(hidden_states)
+
+        # Loss (optional, only when labels are passed)
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+            )
+
+        return CausalLMOutputWithCrossAttentions(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+from transformers import AutoConfig, AutoModelForCausalLM
+
+# Register Aurora config
+AutoConfig.register("aurora-gpt2", AuroraGPT2Config)
+
+# Register Aurora model class
+AutoModelForCausalLM.register(AuroraGPT2Config, AuroraGPT2ForCausalLM)

config.json

@@ -2,7 +2,7 @@
   "_name_or_path": "./small_gpt",
   "activation_function": "gelu_new",
   "architectures": [
-    "GPT2LMHeadModel"
+    "AuroraGPT2ForCausalLM"
   ],
   "attn_pdrop": 0.1,
   "bos_token_id": 50256,
@@ -10,7 +10,7 @@
   "eos_token_id": 50256,
   "initializer_range": 0.02,
   "layer_norm_epsilon": 1e-05,
-  "model_type": "gpt2",
+  "model_type": "aurora-gpt2",
   "n_embd": 256,
   "n_head": 4,
   "n_inner": null,