Update to checkpoint 35000 with fixed Pre-LN architecture (iter 35k, loss 3.46)

config.json

@@ -1,30 +1,22 @@
 {
+  "activation_function": "gelu",
   "architectures": [
-    "GPT2LMHeadModel"
+    "NanoGPTLMHeadModel"
   ],
-  "model_type": "gpt2",
-  "vocab_size": 50257,
-  "n_positions": 512,
+  "attn_pdrop": 0.1,
+  "bos_token_id": 50256,
+  "embd_pdrop": 0.1,
+  "eos_token_id": 50256,
+  "initializer_range": 0.02,
+  "layer_norm_epsilon": 1e-05,
+  "model_type": "nanogpt",
   "n_embd": 384,
-  "n_layer": 8,
   "n_head": 8,
   "n_inner": 1536,
-  "activation_function": "gelu",
+  "n_layer": 8,
+  "n_positions": 512,
   "resid_pdrop": 0.1,
-  "embd_pdrop": 0.1,
-  "attn_pdrop": 0.1,
-  "layer_norm_epsilon": 1e-05,
-  "initializer_range": 0.02,
-  "bos_token_id": 50256,
-  "eos_token_id": 50256,
-  "tie_word_embeddings": true,
   "torch_dtype": "float32",
-  "transformers_version": "4.35.0",
-  "mlx_training": {
-    "framework": "MLX",
-    "iterations": 20000,
-    "final_loss": 0.7582720518112183,
-    "dataset": "tinystories",
-    "max_tokens": 2000000
-  }
-}
+  "transformers_version": "4.51.3",
+  "vocab_size": 50257
+}

generation_config.json

@@ -1,9 +1,6 @@
 {
+  "_from_model_config": true,
   "bos_token_id": 50256,
   "eos_token_id": 50256,
-  "max_length": 512,
-  "temperature": 1.0,
-  "top_k": 50,
-  "top_p": 0.95,
-  "do_sample": true
-}
+  "transformers_version": "4.51.3"
+}

modeling_nanogpt.py

@@ -0,0 +1,267 @@
+"""
+Custom HuggingFace-compatible GPT model with Pre-LN architecture
+Matches the MLX nanoGPT implementation exactly
+"""
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel, PretrainedConfig
+from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions
+import math
+
+
+class NanoGPTConfig(PretrainedConfig):
+    """Configuration for NanoGPT model"""
+    model_type = "nanogpt"
+    
+    def __init__(
+        self,
+        vocab_size=50257,
+        n_positions=512,
+        n_embd=384,
+        n_layer=8,
+        n_head=8,
+        n_inner=1536,
+        activation_function="gelu",
+        resid_pdrop=0.1,
+        embd_pdrop=0.1,
+        attn_pdrop=0.1,
+        layer_norm_epsilon=1e-5,
+        initializer_range=0.02,
+        bos_token_id=50256,
+        eos_token_id=50256,
+        **kwargs
+    ):
+        super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
+        self.vocab_size = vocab_size
+        self.n_positions = n_positions
+        self.n_embd = n_embd
+        self.n_layer = n_layer
+        self.n_head = n_head
+        self.n_inner = n_inner
+        self.activation_function = activation_function
+        self.resid_pdrop = resid_pdrop
+        self.embd_pdrop = embd_pdrop
+        self.attn_pdrop = attn_pdrop
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.initializer_range = initializer_range
+
+
+class NanoGPTAttention(nn.Module):
+    """Multi-head self-attention with Pre-LN"""
+    
+    def __init__(self, config):
+        super().__init__()
+        self.n_head = config.n_head
+        self.n_embd = config.n_embd
+        self.head_dim = self.n_embd // self.n_head
+        self.scale = math.sqrt(self.head_dim)
+        
+        # Combined QKV projection (standard Linear, not Conv1D)
+        self.qkv_proj = nn.Linear(config.n_embd, 3 * config.n_embd)
+        self.out_proj = nn.Linear(config.n_embd, config.n_embd)
+        self.attn_dropout = nn.Dropout(config.attn_pdrop)
+        self.resid_dropout = nn.Dropout(config.resid_pdrop)
+        
+        # Causal mask
+        self.register_buffer(
+            "bias",
+            torch.tril(torch.ones(config.n_positions, config.n_positions)).view(
+                1, 1, config.n_positions, config.n_positions
+            ),
+        )
+        
+    def forward(self, x):
+        B, T, C = x.size()  # batch, sequence length, embedding dim
+        
+        # Project and split into Q, K, V
+        qkv = self.qkv_proj(x)  # (B, T, 3*C)
+        qkv = qkv.view(B, T, 3, self.n_head, self.head_dim)  # (B, T, 3, n_head, head_dim)
+        qkv = qkv.permute(2, 0, 3, 1, 4)  # (3, B, n_head, T, head_dim)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+        
+        # Scaled dot-product attention
+        scores = (q @ k.transpose(-2, -1)) / self.scale  # (B, n_head, T, T)
+        
+        # Apply causal mask
+        scores = scores.masked_fill(self.bias[:, :, :T, :T] == 0, float('-inf'))
+        
+        attn_weights = torch.softmax(scores, dim=-1)
+        attn_weights = self.attn_dropout(attn_weights)
+        
+        # Combine heads
+        out = attn_weights @ v  # (B, n_head, T, head_dim)
+        out = out.transpose(1, 2).contiguous().view(B, T, C)  # (B, T, C)
+        
+        return self.resid_dropout(self.out_proj(out))
+
+
+class NanoGPTMLP(nn.Module):
+    """Feed-forward network"""
+    
+    def __init__(self, config):
+        super().__init__()
+        self.fc1 = nn.Linear(config.n_embd, config.n_inner)
+        self.fc2 = nn.Linear(config.n_inner, config.n_embd)
+        self.act = nn.GELU()
+        self.dropout = nn.Dropout(config.resid_pdrop)
+        
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.dropout(x)
+        x = self.fc2(x)
+        return self.dropout(x)
+
+
+class NanoGPTBlock(nn.Module):
+    """Transformer block with Pre-LN architecture"""
+    
+    def __init__(self, config):
+        super().__init__()
+        self.ln1 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
+        self.attn = NanoGPTAttention(config)
+        self.ln2 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
+        self.mlp = NanoGPTMLP(config)
+        
+    def forward(self, x):
+        # Pre-norm architecture (LayerNorm before attention/MLP)
+        x = x + self.attn(self.ln1(x))
+        x = x + self.mlp(self.ln2(x))
+        return x
+
+
+class NanoGPTModel(PreTrainedModel):
+    """NanoGPT model with Pre-LN architecture"""
+    config_class = NanoGPTConfig
+    
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+        
+        # Embeddings
+        self.token_embedding = nn.Embedding(config.vocab_size, config.n_embd)
+        self.position_embedding = nn.Embedding(config.n_positions, config.n_embd)
+        self.drop = nn.Dropout(config.embd_pdrop)
+        
+        # Transformer blocks
+        self.blocks = nn.ModuleList([NanoGPTBlock(config) for _ in range(config.n_layer)])
+        
+        # Final layer norm
+        self.ln_f = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
+        
+        # LM head (tied with token_embedding)
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        
+        # Tie weights
+        self.lm_head.weight = self.token_embedding.weight
+        
+        # Initialize weights
+        self.apply(self._init_weights)
+        
+    def _init_weights(self, module):
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if isinstance(module, nn.Linear) and module.bias is not None:
+                module.bias.data.zero_()
+    
+    def forward(
+        self,
+        input_ids=None,
+        past_key_values=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        labels=None,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        
+        if input_ids is not None:
+            batch_size, seq_length = input_ids.size()
+        else:
+            batch_size, seq_length = inputs_embeds.size()[:-1]
+        
+        if position_ids is None:
+            if input_ids is not None:
+                device = input_ids.device
+            elif inputs_embeds is not None:
+                device = inputs_embeds.device
+            else:
+                device = next(self.parameters()).device
+            position_ids = torch.arange(seq_length, dtype=torch.long, device=device)
+            position_ids = position_ids.unsqueeze(0).expand(batch_size, -1)
+        
+        # Embeddings
+        if inputs_embeds is None:
+            inputs_embeds = self.token_embedding(input_ids)
+        
+        position_embeds = self.position_embedding(position_ids)
+        hidden_states = self.drop(inputs_embeds + position_embeds)
+        
+        # Transformer blocks
+        for block in self.blocks:
+            hidden_states = block(hidden_states)
+        
+        # Final layer norm
+        hidden_states = self.ln_f(hidden_states)
+        
+        # LM head
+        lm_logits = self.lm_head(hidden_states)
+        
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = lm_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))
+        
+        if not return_dict:
+            output = (lm_logits,)
+            return ((loss,) + output) if loss is not None else output
+        
+        return CausalLMOutputWithCrossAttentions(
+            loss=loss,
+            logits=lm_logits,
+            past_key_values=None,
+            hidden_states=None,
+            attentions=None,
+        )
+
+
+class NanoGPTLMHeadModel(PreTrainedModel):
+    """Causal language model wrapper"""
+    config_class = NanoGPTConfig
+    
+    def __init__(self, config):
+        super().__init__(config)
+        self.transformer = NanoGPTModel(config)
+        
+    def forward(self, *args, **kwargs):
+        return self.transformer(*args, **kwargs)
+    
+    def generate(self, *args, **kwargs):
+        """Use HF's generate method"""
+        # Remove unused kwargs that transformers might pass
+        kwargs.pop("attention_mask", None)
+        kwargs.pop("token_type_ids", None)
+        return super().generate(*args, **kwargs)
+    
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
+        # Remove unused kwargs
+        kwargs.pop("attention_mask", None)
+        kwargs.pop("token_type_ids", None)
+        
+        # Our model doesn't support KV caching, so we need to pass the full sequence
+        return {
+            "input_ids": input_ids,
+        }
+    
+    def can_generate(self):
+        """Indicate this model can generate"""
+        return True

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6d21ad8e8646491e7510cb58cc5d542ca21db6a8e174f5399fba9546662cf317
+size 143190611

special_tokens_map.json

@@ -1,5 +1,23 @@
 {
-  "bos_token": "<|endoftext|>",
-  "eos_token": "<|endoftext|>",
-  "unk_token": "<|endoftext|>"
+  "bos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
 }

tokenizer_config.json

@@ -1,4 +1,5 @@
 {
+  "add_bos_token": false,
   "add_prefix_space": false,
   "added_tokens_decoder": {
     "50256": {
@@ -13,8 +14,10 @@
   "bos_token": "<|endoftext|>",
   "clean_up_tokenization_spaces": false,
   "eos_token": "<|endoftext|>",
+  "errors": "replace",
   "extra_special_tokens": {},
   "model_max_length": 1024,
+  "pad_token": null,
   "tokenizer_class": "GPT2Tokenizer",
   "unk_token": "<|endoftext|>"
 }