nanollm_wiki_20.2M

NanoLLM 20.2M Wiki

nanollm_wiki_20.2M is a lightweight, decoder-only transformer model designed for educational purposes, fast training experimentation, and local generation on Apple Silicon using Apple's MLX framework.

The model is trained on a 100k subset of English Wikipedia articles (chonkie-ai/wikipedia-100k) to demonstrate basic autoregressive text generation capabilities with minimal compute requirements.

Model Details

• Developed by: samairtimer
• Model Type: Decoder-only Transformer (Custom NanoLLM Architecture)
• Language(s): English
• License: MIT
• Framework: Apple MLX
• Tokenizer: GPT-2 (via tiktoken, vocab size 50,257)

Model Architecture & Hyperparameters

Unlike standard Llama or GPT architectures, this model is a highly simplified attention-only transformer that omits the typical Feed-Forward Neural Network (FFN) sub-blocks and layer normalization to keep training extremely fast and lightweight.

Parameter Count Breakdown

The model contains exactly 20,208,000 parameters (~20.2M).

$$ \text{Total Parameters} = \text{Embeddings} + \text{Transformer Blocks} + \text{Output Projection} $$

1. Embedding Layer:

   • Token Embeddings: $50,257 \times 192 = 9,649,344$
   • Positional Embeddings: $128 \times 192 = 24,576$
   • Subtotal: 9,673,920 parameters

2. Transformer Blocks (6 Layers):

   • Each layer contains only a Multi-head Attention block (no biases, no feed-forward projection):
     • Query Projection: $192 \times 192 = 36,864$
     • Key Projection: $192 \times 192 = 36,864$
     • Value Projection: $192 \times 192 = 36,864$
     • Output Projection: $192 \times 192 = 36,864$
     • Subtotal per layer: $147,456$ parameters
   • Total for 6 layers: $147,456 \times 6 = \mathbf{884,736}$ parameters

3. Output Layer (LM Head):

   • Untied weights, linear layer with bias=False: $192 \times 50,257 = \mathbf{9,649,344}$ parameters

$$\text{Grand Total} = 9,673,920 + 884,736 + 9,649,344 = 20,208,000 \text{ parameters}$$

---

Training Recipe

Dataset

• Name: chonkie-ai/wikipedia-100k
• Size: 100,000 high-quality parsed English Wikipedia samples (~449 MB of raw text).
• Format: Combined articles delimited by the <|endoftext|> token.

Training Configuration

• Optimizer: AdamW (learning_rate=3e-4)
• Batch Size: 32
• Sequence Length: 128 tokens
• Epochs: 5 epochs (~15,600 steps total)
• Precision: float32 (saving to .safetensors size of ~80.8 MB)
• Hardware: Apple Silicon Mac GPU (accelerated via MLX Metal backend)

Loss Curve

• Initial Loss: 10.8269
• End of Epoch 1 (Step 3100): 5.3371
• End of Epoch 2 (Step 6200): 5.0588
• End of Epoch 3 (Step 9350): 4.8099
• End of Epoch 4 (Step 12450): 4.4191
• End of Epoch 5 (Step 15600): 4.3014

---

How to Use

To use this model for text generation, define the architecture matching the parameters and load the saved .safetensors weights using MLX.

1. Requirements

Install the dependencies:

pip install mlx tiktoken huggingface_hub

2. Inference Code

import os
import mlx.core as mx
import mlx.nn as nn
import tiktoken
from huggingface_hub import hf_hub_download

# Define model architecture identical to training
class TokenAndPositionEmbedding(nn.Module):
    def __init__(self, maxlen: int, vocab_size: int, embed_dim: int):
        super().__init__()
        self.token_emb = nn.Embedding(vocab_size, embed_dim)
        self.pos_emb = nn.Embedding(maxlen, embed_dim)

    def __call__(self, x):
        seq_len = x.shape[1]
        positions = mx.arange(seq_len)[None, :]
        return self.token_emb(x) + self.pos_emb(positions)

class TransformerBlock(nn.Module):
    def __init__(self, emed_dim: int, num_heads: int, ff_dim: int):
        super().__init__()
        self.attention = nn.MultiHeadAttention(emed_dim, num_heads)
        
    def __call__(self, x, mask=None):
        attn_out = self.attention(x, x, x, mask=mask)
        return x + attn_out

class NanoLLM(nn.Module):
    def __init__(self, maxlen: int, vocab_size: int, embed_dim: int, num_heads: int, feed_forward_dim: int, num_transformer_blocks: int):
        super().__init__()
        self.maxlen = maxlen
        self.embedding = TokenAndPositionEmbedding(maxlen, vocab_size, embed_dim)
        self.transformer_blocks = [
            TransformerBlock(embed_dim, num_heads, feed_forward_dim)
            for _ in range(num_transformer_blocks)
        ]
        self.output_layer = nn.Linear(embed_dim, vocab_size, bias=False)

    def __call__(self, token_ids):
        seq_len = token_ids.shape[1]
        mask = nn.MultiHeadAttention.create_additive_causal_mask(seq_len)
        x = self.embedding(token_ids)
        for block in self.transformer_blocks:
            x = block(x, mask=mask)
        return self.output_layer(x)

# 1. Download weights from Hugging Face
repo_id = "samairtimer/nanollm_wiki_20.2M"
weights_path = hf_hub_download(repo_id=repo_id, filename="wikipedia_checkpoint.safetensors")

# 2. Initialize and load weights
tokenizer = tiktoken.get_encoding("gpt2")
model = NanoLLM(
    maxlen=128,
    vocab_size=tokenizer.n_vocab,
    embed_dim=192,
    num_heads=6,
    feed_forward_dim=512,
    num_transformer_blocks=6
)
model.load_weights(weights_path)
print("Model loaded successfully!")

# 3. Autoregressive Generation function
def generate(model, tokenizer, prompt, max_new_tokens=100, temperature=0.6):
    tokens = tokenizer.encode(prompt)
    x = mx.array(tokens)[None, :]
    end_token_id = tokenizer.encode('<|endoftext|>', allowed_special={'<|endoftext|>'})[0]
    
    print(prompt, end="", flush=True)
    for _ in range(max_new_tokens):
        if x.shape[1] > model.maxlen:
            x = x[:, -model.maxlen:]
        
        logits = model(x)
        next_token_logits = logits[0, -1, :]
        
        if temperature == 0.0:
            next_token = mx.argmax(next_token_logits, axis=-1).item()
        else:
            next_token = mx.random.categorical(next_token_logits / temperature, axis=-1).item()
            
        if next_token == end_token_id:
            break
            
        word = tokenizer.decode([next_token])
        print(word, end="", flush=True)
        x = mx.concatenate([x, mx.array([[next_token]])], axis=1)
    print("\n")

# Run generation
generate(model, tokenizer, "Cats are", max_new_tokens=50, temperature=0.4)

---

Limitations & Biases

• Extremely Small Scale: At 20.2M parameters, the model is not capable of complex reasoning, logical deduction, or factual lookup.
• Attention-Only: The omission of MLP (Feed-Forward) layers limits the key-value lookup and relational storage capacity of the network, resulting in higher recurrence of repetitive structures.
• Short Context Window: The 128-token context window restricts its capacity to maintain long-term coherence.
• Factual Inaccuracies: The model is trained on a small subset of Wikipedia (100k entries) for 5 epochs. Any generated output should be treated strictly as synthetic text generation rather than historical or scientific fact.