TinyLLM — End-to-End Training & Alignment

Why build one from scratch

Using a language model and understanding one are different things. This project implements the full stack — tokenizer, architecture, pretraining loop, then the alignment stages that turn a next-token predictor into something that follows instructions — on a scale that fits a single GPU.

Architecture choices

The model is decoder-only, following the GPT family, but uses the architectural refinements that have become standard since the original design rather than reproducing it verbatim.

• Rotary position embeddings (RoPE) encode position by rotating query and key vectors, which extrapolates to longer sequences better than learned absolute positions.
• Grouped-query attention (GQA) shares key and value heads across groups of query heads, cutting the memory bandwidth cost of the attention cache with little quality loss.
• Pre-LayerNorm places normalization before each sublayer rather than after, which makes deep transformers substantially more stable to train.
• GELU activations and WordPiece tokenization complete the stack.

Training and alignment

Pretraining is standard next-token language modelling with mixed precision and AdamW under a warmup schedule — the combination that makes single-GPU training of a model this size tractable.

Alignment is done in two stages. Instruction tuning adapts the pretrained model to follow prompts. Preference ranking then refines behaviour using a direct preference objective rather than full reinforcement learning from human feedback — the same goal as RLHF, without the separate reward model and policy optimization loop, which is a meaningful simplification at this scale.