docs

APEX-1 Documentation

Complete beginner-to-expert guide — learn how a real AI language model is designed, coded, and trained by reading through every file in this project.

---

📚 How to Read This Documentation

If you are new to AI/ML, start at 00-introduction.md and read in order. Each document builds on the previous one. By the end, you will understand how a modern large language model is built from scratch.

If you want to jump to a specific topic, use the table below.

---

Table of Contents

🟢 Part 1 — Foundations (Start Here)

#	File	What You Will Learn
00	Introduction — What is an LLM?	What AI models are, why we build them
01	Project Structure	Every file and folder explained
02	Configuration	The blueprint (config.py)
03	Tokenizer	Turning text into numbers

🔵 Part 2 — Building Blocks

#	File	What You Will Learn
04	Embeddings & RMSNorm	Word vectors, normalisation
05	Positional Encoding — RoPE & YaRN	Teaching position to the model
06	Attention Masks	Who can see what

🟣 Part 3 — Attention Mechanisms

#	File	What You Will Learn
07	Multi-Head Latent Attention (MLA)	93% KV cache reduction
08	Grouped Query Attention + Sliding Window	Efficient local attention

🟠 Part 4 — Feed-Forward Networks & Experts

#	File	What You Will Learn
09	FFN & SwiGLU	The fact-memory layer
10	Mixture of Experts (MoE)	256 specialists, 2 active
11	Dynamic Skip Gate	Skipping easy tokens
12	Auxiliary-Loss-Free Load Balancer	Balancing experts
13	Multi-Token Prediction Head	Predicting 4 tokens at once

🔴 Part 5 — The Full Model

#	File	What You Will Learn
14	Transformer Block	One complete layer
15	APEX1Model — The Complete Model	All layers assembled

🟡 Part 6 — Training

#	File	What You Will Learn
16	Training Losses	How the model learns
17	Optimizer & LR Scheduler	AdamW + cosine warmup
18	Training Pipeline	The full training loop
19	Checkpointing	Saving and restoring
20	Datasets	Loading and preparing data

⚪ Part 7 — Text Generation

#	File	What You Will Learn
21	Sampling Strategies	How text is generated
22	Speculative Decoding	3× faster generation
23	Thinking Mode	Built-in reasoning scratchpad

🟤 Part 8 — Alignment & Safety

#	File	What You Will Learn
24	Reward Model	What humans prefer
25	DPO — Direct Preference Optimization	Training on preferences
26	GRPO — Group Relative Policy Optimization	RL without a value function
27	Process Reward Model	Rewarding good reasoning
28	Constitutional AI	Safety baked in
29	Combined Reward	All signals together

⚫ Part 9 — Utilities & Walkthrough

#	File	What You Will Learn
30	Utilities	Shape checker, FLOPs, param counter
31	End-to-End Walkthrough	Full journey, runnable code

---

Mathematical Reference (Advanced)

The original mathematical reference documents are still available:

Part	Topics
Math Ref Part 1	Embedding, RMSNorm, RoPE, YaRN
Math Ref Part 2	SDPA, MHA, GQA, MLA, Masks
Math Ref Part 3	SwiGLU, MoE, Load Balancing, Skip Gate
Math Ref Part 4	AdamW, LR Schedule, DPO, GRPO, Sampling

---

Full Architecture Document

• APEX-1 Model Architecture — Complete technical design

---

Quick Start (Run the Code Now)

# 1. Install
pip install -e ".[all]"

# 2. Run a forward pass
python examples/forward_pass_demo.py

# 3. Generate text
python examples/generation_demo.py

# 4. Generate with thinking mode
python examples/thinking_mode_demo.py

# 5. Run all tests
pytest tests/ -v