PREE

Built a decoder-only Transformer from scratch in pure PyTorch.

Overview

PREE is a fully custom decoder-only Transformer architecture implemented from scratch. It supports the完整的 transformer pipeline from tokenization through training to inference, with performance optimizations targeting both memory efficiency and training speed.

Problem

Most transformer implementations rely on bloated frameworks and abstract away the core mechanics. Understanding transformer internals requires building one. PREE was created to demystify every component — from tokenization to weight tying — and provide a clear, readable, production-conscious implementation.

Architecture

The model follows the GPT-2 family design with several key modifications:

Decoder-only autoregressive architecture

Grouped Query Attention (GQA) for efficient inference with reduced KV-cache memory

Rotary Position Embeddings (RoPE) for superior length generalization

Pre-norm layer normalization for training stability

Tied embeddings between the input token embedding and output projection head

`

Input Tokens → BPE Tokenizer → Embedding + RoPE → N × [GQA Block] → LM Head → Logits

`

Each GQA block contains:

Grouped Query Attention with configurable head groups

SwiGLU feed-forward network

RMSNorm pre-normalization

Residual connections

Pipeline

Training Pipeline

Data loading: Streaming dataset loader from disk with memory-mapped files

Tokenization: Custom BPE tokenizer trained on the target corpus

Batching: Dynamic batch sizing based on sequence length to maximize GPU utilization

Mixed precision: Forward pass in FP16, master weights in FP32

Gradient accumulation: Simulates larger batch sizes across multiple micro-batches

Optimizer step: 8-bit AdamW with decoupled weight decay

Gradient checkpointing: Trades compute for memory by recomputing activations during backward pass

Learning rate scheduling: Warmup followed by cosine decay

Inference Pipeline

KV-cache: Persistent key-value cache across decode steps for GQA

Speculative decoding: Draft model proposes tokens, target model verifies in parallel

Quantization-aware: INT8 quantization support for weights and activations

Training

Parameter	Value
-----------	-------
Vocabulary Size	50,304
Context Length	2,048
Hidden Dimension	768
Layers	12
Attention Heads	12 (4 KV groups)
Head Dimension	64
Training Tokens	1B+
Batch Size	512 sequences
Learning Rate	3e-4 (warmup) → 3e-5 (cosine decay)
Optimizer	8-bit AdamW (lr=3e-4, wd=0.1)
Precision	FP16 (forward), FP32 (master)
Gradient Checkpointing	Enabled
Hardware	Single NVIDIA A100 80GB
Training Time	~48 hours

Challenges

Numerical Stability

FP16 training introduced underflow issues in softmax computation. Solution: applied scaling before the exp operation and used dynamic loss scaling.

Memory Bottleneck

Standard attention with full KV cache consumed 40+ GB for long sequences. Switched to GQA with 4 KV groups, reducing memory by 4x. Combined with gradient checkpointing, full 2048-sequence training fits in 80GB.

Custom BPE Training

Training a BPE tokenizer from scratch required careful handling of merges, vocabulary sizing, and handling of special tokens. The tokenizer achieves near-optimal compression on the training corpus.

RoPE Implementation

Implementing RoPE correctly required understanding the rotation matrix formulation and handling the caching of cos/sin tables efficiently. Final implementation uses precomputed cached tables for all positions up to the context length.

Optimization

8-bit AdamW

Used bitsandbytes for 8-bit optimizer states, reducing memory footprint by 3x compared to FP32 Adam while maintaining convergence quality.

Gradient Checkpointing

Checkpointed activations at each transformer layer instead of storing all intermediate states. This increases compute per step by ~20% but reduces memory by ~60%.

Flash Attention

Integrated Flash Attention 2 for the GQA blocks, achieving 2-3x speedup on the attention computation with no numerical accuracy loss.

Kernel Fusion

Fused the SwiGLU activation with the linear projection, reducing memory bandwidth overhead by eliminating intermediate tensor materialization.

Results

Metric	Value
--------	-------
Final Loss	2.83
Validation Loss	2.91
Perplexity	18.5
Tokens/sec (forward)	12,400
Tokens/sec (decode)	3,200
Memory Usage (peak)	62 GB

The model demonstrates strong few-shot capabilities on standard benchmarks, outperforming similarly-sized models trained without GQA or RoPE.

Lessons Learned

GQA is a drop-in efficiency win: 4 KV groups achieve near-MHA quality at half the KV-cache memory.

RoPE generalizes better than learned positional embeddings: Especially for sequence lengths not seen during training.

Gradient checkpointing is essential: Without it, training at 2048 context length is infeasible on a single 80GB GPU.

Mixed precision requires care: Loss scaling and dynamic casting are critical to avoid nan gradients.

BPE tokenizer quality matters: A poorly trained tokenizer degrades downstream model performance significantly.

Future Work

Implement Mixture of Experts (MoE) routing for conditional computation

Add multi-head latent attention for even more efficient KV-cache usage

Explore block-sparse attention patterns for longer contexts

Build an ONNX export path for deployment to edge devices

Add controllable generation via classifier-free guidance

Source Code

Github Repository