Tiny Shakespeare GPT (MoE)
This is a small custom Generative Pre-trained Transformer (GPT) model trained on the Tiny Shakespeare dataset. It features a Mixture of Experts (MoE) architecture and Grouped-Query Attention (GQA).
Model Details
- Architecture: GPT with Mixture of Experts (MoE)
- Total Parameters: 32,101,504
- Block Size (Context Window): 256 tokens
- Layers: 4
- Heads: 4
- KV Heads: 2 (Grouped-Query Attention)
- Embedding Size: 128
- Total Experts (MoE): 8
- Experts Per Token: 2
Training Details
- Dataset: Tiny Shakespeare
- Total Iterations: 500
- Final Train Loss: 3.6204
- Final Validation Loss: 4.8322 (Best: 4.7383)
Hyperparameters
- Batch Size: 3
- Max Iters: 500
- Learning Rate: 0.001 (Min LR: 0.0001)
- Warmup Iters: 100
- Weight Decay: 0.1
- Gradient Accumulation Steps: 16
Usage
You can use this model by cloning the original repository, as it requires the custom architecture implementation.
import torch
from src.model import GPT, GPTConfig
# Load the configuration and initialize the model
config = GPTConfig(
block_size=256,
n_layer=4,
n_head=4,
n_kv_head=2,
n_embd=128,
n_experts=8,
num_experts_per_tok=2
)
model = GPT(config)
# Load the weights from the safetensors file
from safetensors.torch import load_file
state_dict = load_file("model.safetensors")
model.load_state_dict(state_dict)
model.eval()
model.to('cuda' if torch.cuda.is_available() else 'cpu')
# You can now generate text using the `model.generate` method
# (assuming you have the tokenizer loaded from the original repo)
Loss Curve
