Transformer From Scratch — 123M Parameter Decoder-Only Language Model
A decoder-only transformer built entirely from first principles in PyTorch — every component (embeddings, RoPE, multi-head attention, RMSNorm, SwiGLU, output projection) implemented and numerically verified against reference math, then trained to convergence on WikiText-103. Accelerated with a hand-written, hand-derived CUDA attention kernel, verified correct to float32 precision.
Model Details
| |
|---|
| Parameters | 123,551,232 (123.6M) |
| Architecture | Decoder-only, RoPE, pre-norm residuals, SwiGLU, tied embeddings |
| Layers | 12 |
| Hidden dimension | 768 |
| Attention heads | 12 |
| Feed-forward hidden dim | 2048 (SwiGLU-gated) |
| Vocabulary | 50,257 (GPT-2 tokenizer) |
| Max sequence length | 1024 |
| Training data | WikiText-103 |
| Training steps | 36,621 (~3 epochs) |
| Hardware | NVIDIA A40, ~9 hours |
Evaluation Results
Evaluated on WikiText-103's held-out validation split (~123K tokens, 30 batches):
| Metric | Value |
|---|
| Validation perplexity | 19.24 |
| Next-token accuracy (top-1) | 44.61% |
| Next-token accuracy (top-5) | 66.07% |
| Train/val perplexity gap | +0.30 (minimal overfitting) |
For context: random-chance accuracy on a 50,257-word vocabulary is ~0.002%. A next-token accuracy of 44.61% is consistent with published results for GPT-2-class (124M parameter) models on comparable held-out text.
Training Curve
Validation perplexity decreased from 277.78 (step 1000) to 19.24 (step 36000), with a clear plateau in the final third of training.
Sample Generation
"The history of the game is not known; it was most recently held in a building dating from 1875, when the city of Liverpool was founded. Liverpool played at Anfield, and Liverpool played at the current home of the Liverpool F.C.'s, until the 1970s."
"The film was released in the United States and Canada on April 9 and in the UK on June 27. It grossed $4,083,854 in the US and $6,050,989 worldwide."
Custom CUDA Attention Kernel
A hand-derived forward and backward CUDA kernel replaces the attention computation (softmax(QKᵀ/√d)V), bridged into PyTorch via torch.autograd.Function.
| Result |
|---|
| Forward correctness vs. PyTorch reference | max error 1.49e-07 |
| Backward correctness (all gradients) | max error < 8e-06 |
| Speed (forward+backward) | 13.7% of PyTorch's native throughput |
The kernel is a
correctness-first, unoptimized implementation (single-thread reductions, atomic-based accumulation, no shared-memory tiling) — the speed gap is expected and explainable, not a defect. See the
GitHub repo for the full derivation and benchmark methodology. The trained model above uses standard PyTorch attention; the CUDA kernel's correctness and trainability were verified independently.
How to Use
1from huggingface_hub import snapshot_download
2import sys
3
4model_dir = snapshot_download(repo_id="krishang-raju/transformer-scratch-123M-wikitext")
5sys.path.insert(0, model_dir)
6
7from modeling_transformer import TransformerModel
8
9model = TransformerModel.from_pretrained(model_dir, device="cpu")
10
11from transformers import GPT2Tokenizer
12tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
13
14prompt = "The history of"
15input_ids = tokenizer.encode(prompt, return_tensors="pt")
16output_ids = model.generate(input_ids, max_new_tokens=60)
17print(tokenizer.decode(output_ids[0]))
Limitations
- Trained on WikiText-103 (encyclopedic register) — will not generalize well to conversational, code, or highly domain-specific text.
- No fp16/bf16 support in the custom CUDA kernel (float32 only).
- Small-model repetition artifacts are present in longer generations (see distinct-n scores in the full evaluation report on GitHub).
- Evaluated on a ~123K token sample of the validation split, not the full set.
Citation
If referencing this project, please link to the
GitHub repository, which contains the full phase-by-phase verification methodology, training logs, and CUDA kernel derivation.
