README.md

---
language: en
license: mit
tags:
  - gpt
  - pytorch
  - from-scratch
  - course-project
  - decoder-only
---

# 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

```python
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.