Mini-GPT (mini-gpt2_decoder)

A small GPT-style decoder-only transformer, trained from scratch.

Model Details

Architecture	Decoder-only transformer (GPT-2 style, pre-norm)
Parameters	~30.0M
Layers	6
Attention heads	6
Embedding dimension	384
Context length	256
Vocabulary	50257 tokens (GPT-2 BPE, via tiktoken)

Training Data

A ~100MB slice of OpenWebText, tokenized with GPT-2's BPE tokenizer.

Project Context

This model was trained and optimized across six phases:

1. FP32 baseline training
2. Mixed precision (AMP)
3. Memory engineering (gradient checkpointing + gradient accumulation)
4. A custom fused attention kernel written in Triton
5. KV-cached autoregressive generation
6. INT8 post-training quantization (evaluated separately, CPU-only -- this upload contains the FP32 weights, see the scope note in upload_to_hf.py in the project repo for why)

Optimization Results

Mini-GPT Optimization Project -- Final Summary

Phases 1-3: Training Optimizations

• baseline: 1170.6s, peak memory 8208MB, final val loss 5.5229
• mixed_precision: 484.0s, peak memory 8208MB, final val loss 5.5225
• mixed_precision: 483.7s, peak memory 8208MB, final val loss 5.5224
• memory_engineering: 1763.6s, peak memory 8208MB, final val loss 5.0966
• From baseline to memory_engineering: 0.66x time change, 1.00x memory change, validation loss moved from 5.5229 to 5.0966 (should stay close -- these are meant to be near-lossless optimizations, not accuracy tradeoffs).

Phase 4: Custom Fused Attention Kernel

• Average forward-pass speedup across tested sequence lengths: 0.87x (vanilla SDPA vs. our Triton kernel).
• Best result: 0.89x at seq_len=64 (18.006ms -> 20.132ms).
• Scope: forward pass only -- see model/attention.py's docstring for the autograd/backward-pass caveat.

Phase 5: KV-Cache Generation

• Speedup grew from 1.10x at 20 tokens to 1.23x at 200 tokens -- confirming the expected O(n) vs O(n^2) behavior (naive generation's wasted recomputation grows with length; cached generation's doesn't).
• Verified mathematically exact: token-for-token identical output under greedy decoding, not an approximation.

Phase 6: INT8 Quantization

• Model size: 114.6MB -> 102.7MB (1.12x smaller).
• CPU latency: 180.14ms -> 165.93ms (1.09x faster).
• Perplexity cost: 163.00 -> 156.69 (+-6.31) -- the one deliberate accuracy-for-efficiency trade in the whole project.

Overall Story

Four of the five speed/memory optimizations (mixed precision, gradient checkpointing, the fused kernel, KV-cache) are effectively free -- same output, faster or leaner. Only quantization (Phase 6) trades a small, explicitly measured amount of accuracy for a large gain in deployability on cheap hardware. That asymmetry -- most optimizations are free, one isn't, and we can say exactly how much it costs -- is the actual finding of this project.

Usage

import json
import torch
from safetensors.torch import load_model

from model.gpt import GPTConfig, MiniGPT   # from this project's repo

with open("config.json") as f:
    config_dict = json.load(f)
config = GPTConfig(**config_dict)

model = MiniGPT(config)
load_model(model, "model.safetensors")
model.eval()

# Naive generation (see generate.py in the project repo for the KV-cached,
# much faster version used during actual benchmarking)
idx = torch.zeros((1, 1), dtype=torch.long)  # or your own prompt token ids
out = model.generate(idx, max_new_tokens=50)

Limitations

Trained on a relatively small (~100MB) data slice as a course project focused on systems optimization, not on maximizing language modeling quality -- expect grammatically plausible but not highly coherent output over long spans.