Repository navigation
v0.1-alpha: "Full stack" LLM training prototype
📦 v0.1-alpha — First Full Stack Prototype
🚀 Overview
This release marks the first working prototype of a minimal DeepSeek-style LLM training stack.
✅ Full model initialization
✅ Unit and integration tests passing
✅ Training loop executes and loss is finite
✅ CPU and CUDA support (single GPU)
✨ Features
-
Model Architecture
- DeepSeekTransformer backbone
- Multi-Head Self Attention (MHSA) with Rotary Position Embeddings (RoPE)
- MoE (Mixture of Experts) layer support
- Causal Language Modeling (LM) head
-
Training Loop
- Supports gradient accumulation
- Autocast (
torch.amp) support for mixed precision - Gradient clipping support
- Checkpoint saving every N steps
- Validation loss evaluation during training
- Optimizer and scheduler (AdamW + one-cycle-LR) setup
-
Caching / Efficiency
- KVCache implemented for fast autoregressive decoding
- Efficient
.generate()sampling loop for small batch generation
-
Testing
- ✅ Unit tests for model components (RoPE, MHSA, Transformer Block, MoE)
- ✅ Integration tests for multi-step training
- ✅ Generation shape and loss sanity checks
⚠️ Known Limitations
- No multi-GPU (
torch.distributed) support yet. - MoE expert routing is basic; scaling behavior untested.
- Tokenizer / dataset ingestion is synthetic or minimal.
- No extensive benchmarking (training throughput, memory).
- Limited hyperparameter flexibility (e.g., fixed dropout schedules).
🛠 How to Use
# Install dependencies
uv pip install -e .
# Run unit + integration tests
uv run pytest
# Run a small pretraining loop
uv run python3 train/pretrain.py --hf-dataset-name='wikipedia'📅 Release Motivation
This v0.1-alpha release represents a first major engineering milestone:
- From scratch LLM training infrastructure ✅
- Basic reproducibility ✅
- Sets the foundation for future scaling and research experimentation ✅
--
In subsequent iterations, we plan to start by creating a very small experiment to ensure efficient learning and correct model architecture. Once a 'baseline' is established, we can focus on cleaning up old code or expanding to new features. Only then shall we consider scaling up to a larger scale experiment (since if we skip any of the aforementioned steps, it could be possible to introduce a meaningful and non-trivial regression into the codebase without knowing).