The goal of LiteGPT-16M is not to achieve state-of-the-art performance, but to provide a clean and understandable implementation of a GPT-style language model that can be trained from scratch and extended with modern techniques in future experiments.
1Input Tokens [B, T]
2 │
3 ▼
4┌─────────────────────┐
5│ Token Embeddings │
6│ [vocab, d_model] │
7└─────────────────────┘
8 │
9 ├──────────────┐
10 ▼ │
11┌─────────────────────┐│
12│ Position Embeddings ││
13│ [seq_len, d_model] ││
14└─────────────────────┘│
15 │ │
16 └──────┬───────┘
17 ▼
18 x = tok + pos
19 │
20 ▼
21╔══════════════════════════════╗
22║ Transformer Block × 4 ║
23║ ║
24║ LayerNorm ║
25║ │ ║
26║ ▼ ║
27║ Multi-Head Attention ║
28║ │ ║
29║ ▼ ║
30║ Residual Add ║
31║ │ ║
32║ ▼ ║
33║ LayerNorm ║
34║ │ ║
35║ ▼ ║
36║ FFN ║
37║ │ ║
38║ ▼ ║
39║ Residual Add ║
40╚══════════════════════════════╝
41 │
42 ▼
43┌─────────────────────┐
44│ Final LayerNorm │
45└─────────────────────┘
46 │
47 ▼
48┌─────────────────────┐
49│ LM Head │
50└─────────────────────┘
51 │
52 ▼
53 Logits [B,T,V]
This model is intentionally kept as close to GPT-2 as possible to build a strong understanding of decoder-only transformers before introducing modern architectural improvements.
Features such as RoPE, GQA, FlashAttention, SwiGLU, RMSNorm, and Mixture-of-Experts are intentionally omitted. While these improve efficiency or performance, they add implementation complexity and make it harder to study the core transformer architecture.
The model is designed to train on a single NVIDIA T4 GPU using Google Colab. Model size, context length, and batch size are chosen to fit within limited compute resources.
1shakespeare.txt
2 │
3 ▼
4┌──────────────────┐
5│ GPT-2 Tokenizer │
6│ (tiktoken) │
7└────────┬─────────┘
8 │
9 ▼
10 Token IDs
11 │
12 ▼
13┌──────────────────┐
14│ 90/10 Split │
15│ Train / Val │
16└────────┬─────────┘
17 │
18 ┌────┴────┐
19 ▼ ▼
20train.bin val.bin
21(uint16) (uint16)
