File size: 3,925 Bytes
792b2fb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | ---
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.
|