Savana RDT Model Trainer

A Recurrent-Depth Transformer (RDT) small language model for local training and experimentation on a single consumer GPU. Savana is built with PyTorch, uses uv for dependency management, and drives a Rich TUI for chat, training, and resume workflows. Tokenisation runs through the GPT-2 BPE tokeniser via transformers. One CLI (uv run savana) and one settings file (app/config.py) cover the full pretrain, SFT, and chat pipeline across eight registered size variants.

---

Features

• Recurrent-Depth architecture: one shared transformer block iterated T times per forward pass, with input re-injected every step; sample at higher loop counts than trained on to see depth extrapolation.
• Eight registered sizes: 7m, 15m, 75m, 100m, 250m, 500m, 750m, 1b, defined as ModelConfig presets in app/model_service/variants/registry.py.
• Interactive Rich TUI: menu-driven Chat, Train, Resume, and Quit screens with a numbered train wizard and a live training panel.
• Headless CLI: prep, prep-sft, train, sft, chat, info, variants for scripting and reproducible runs.
• Pretrain plus SFT pipeline: Wikipedia (default) or FineWeb-Edu for pretraining, Alpaca for supervised fine-tuning.
• Centralised configuration: every tunable value is a pydantic-settings field in app/config.py, overridable via .env or SAVANA_* environment variables.
• Mixture-of-Experts blocks: routed and shared experts, per-loop LoRA adapters, and ACT halting on top of GQA or MLA attention.
• Slash-command REPL: adjust loop depth, temperature, top-K, and token budget on the fly during chat.
• Editable install with CUDA 12.4 wheels: uv sync installs the GPU build of PyTorch directly.

---

Project Structure

Savana/
├── .venv/                              # Python virtual environment (auto-created by uv)
├── app/
│   ├── __init__.py
│   ├── __main__.py
│   ├── main.py                         # CLI entry point (registered as `savana`)
│   ├── config.py                       # centralised pydantic-settings
│   ├── interface_service/              # Rich TUI (menus, chat, train wizard)
│   ├── model_service/                  # architecture, tokenisation, variants
│   └── training_service/               # pretraining, SFT, dataset prep
├── docs/                               # technical docs (01_overview, 02_install, ...)
├── models/                             # trained .pt files, one folder per size
├── data/                               # tokenised corpora
├── pyproject.toml                      # project metadata, deps, and CLI script
├── requirements.txt                    # uv-exported deps for cross-compatibility
├── README.md                           # project documentation
└── uv.lock                             # lockfile for reproducible resolution

---

Prerequisites

• Python 3.10+ (uv will pick a compatible interpreter automatically).
• CUDA-capable GPU recommended. CPU works but training is impractically slow.
• uv installed and on PATH.
• Disk space for the tokenised corpora and checkpoints. A 75m run comfortably fits on a modern SSD; 250m and above want tens of gigabytes.
• VRAM budget matched to the chosen variant. 7m and 15m fit on almost anything; 75m and 100m target an 8 GB card; 250m is tight on 8 GB; 500m needs 16 GB+; 750m and 1b want 24 GB+.

---

Getting Started

1. Clone the Project

Clone or download the Savana repository to a folder of your choice.

2. Install Dependencies

From the project root:

uv sync

This will:

• Install PyTorch with CUDA 12.4 (routed via [tool.uv.sources] in pyproject.toml).
• Install runtime deps: transformers, datasets, pydantic-settings, loguru, rich.
• Register the savana console script via [project.scripts].
• Install dev tools (pytest, ruff, black) via the PEP 735 dev group.
• Install Savana itself in editable mode.

3. Verify the Install

Check that CUDA was detected and the CLI resolves:

uv run python -c "import torch; print('cuda:', torch.cuda.is_available())"
uv run savana info

The first should print cuda: True. The second prints the resolved size, paths, and step targets from app/config.py.

4. Launch the TUI

The interactive Rich TUI is the default entry point:

uv run savana             # opens the TUI menu (Chat, Train, Resume, Quit)
uv run savana run         # same thing, explicit
uv run savana chat        # jump straight to the chat screen

From inside the venv you can drop the uv run prefix and just type savana.

5. Run the Pipeline

Headless commands for scripting and reproducible runs:

uv run savana prep --size 75m         # tokenise Wikipedia
uv run savana prep-sft --size 75m     # tokenise Alpaca
uv run savana train --size 75m        # pretrain
uv run savana sft --size 75m          # supervised fine-tune
uv run savana chat --size 75m         # talk to it
uv run savana info                    # print resolved config
uv run savana variants                # list registered sizes

Checkpoints save to models/{size}/pretrain/ and models/{size}/sft/. Sizes coexist on disk; switch by changing --size or the SAVANA_SIZE env var.

---

Screenshots

Main Menu

Rich TUI landing screen with Chat, Train, Resume, and Quit.

Training Menu

The training entry point reached from the main menu, offering pretraining and supervised fine-tuning.

Training Setup

Numbered selection for size, corpus (wikipedia or fineweb_edu), and step count, with a confirmation panel before launch.

---

How It Works

1. Dependency Management

uv handles the virtual environment, PyPI resolution, and the extra PyTorch CUDA 12.4 index. uv sync reads pyproject.toml and uv.lock and produces a reproducible .venv/. The savana console script is registered via [project.scripts], so it resolves inside the venv without a manual pip install.

2. Configuration

Every tunable value lives in app/config.py as a pydantic-settings field. Override any field via environment variable or a .env file. Names follow SAVANA_<FIELD> for top-level values and SAVANA_<GROUP>__<FIELD> for nested groups, e.g. SAVANA_TRAIN__SEQ_LEN=1024. See app/config.py and app/models/config.py for the full list of fields with their defaults.

3. Model and Tokeniser

app/model_service/service.py is the public facade. It builds a model from a ModelConfig preset in app/model_service/variants/registry.py and returns a GPT-2 BPE tokeniser wired through transformers. The ModelConfig selects attention type (gqa or mla), MoE shape (n_experts, n_shared_experts, expert_dim), loop budget (max_loop_iters), and LoRA rank for the per-loop adapters.

4. Training Pipeline

app/training_service/service.py orchestrates two loops: pretrain.py on Wikipedia or FineWeb-Edu, and sft.py on Alpaca. Dataset prep tokenises the corpus once and writes it under data/. The pretrain loop reads fixed-length windows via batching/windows.py; SFT uses the chat-template collator in batching/collate.py. Loss is next-token cross-entropy with a one-position shift. The optimiser runs a warmup plus cosine schedule from optimization/schedule.py, and checkpoints go through checkpointing/store.py so a run can be resumed cleanly.

5. Interactive Interface

app/interface_service/service.py drives the Rich TUI. views/main_menu.py presents the four-way menu; views/train_setup.py is the numbered wizard; views/training_run.py streams live metrics; views/chat.py runs the REPL. Slash commands are parsed in commands/parser.py and sampling is done by generation/sampler.py.

6. Chat Slash Commands

In the chat REPL, slash commands tune sampling on the fly:

Command	Meaning
/loops 12	Recurrent loop depth at inference.
/temp 0.5	Sampling temperature.
/topk 50	Restrict to top-K logits.
/tokens 200	Max new tokens per turn.
/raw	Toggle the chat template.
/info	Print current settings.
/quit	Exit.

7. Architecture in One Paragraph

Input tokens pass through a short Prelude of standard transformer blocks. The output becomes both the initial hidden state h and the input signal e that is re-injected on every loop. The Recurrent Block runs one transformer block T times. At each step h_{t+1} = A·h_t + B·e + Transformer(h_t, e), where A is a diagonal matrix parameterised so spectral radius is below 1 by construction. A learned per-position halting probability lets easy tokens exit the loop early while hard tokens keep iterating. The output is a halting-weighted sum across iterations, then passed through a short Coda and the LM head.

---

Sizes and Time Budgets

Eight variants are registered in app/model_service/variants/registry.py.

Size	Active params	Fits in	Use for
7m	~7M	<2 GB	smoke tests, architecture debugging
15m	~15M	2 GB	CPU-runnable sanity checks
75m	~75M	6-8 GB	local default on a 4070-class GPU
100m	~100M	8 GB	step up from 75m on the same GPU class
250m	~250M	7-8 GB tight	multi-day local run
500m	~500M	16 GB+	mid-range single-GPU training
750m	~750M	24 GB+	high-end single-GPU training
1b	~1.1B	24 GB+	MLA attention, single high-end GPU

Time on a 4070 for the full pretrain plus Alpaca SFT pipeline:

Size	Pretrain	SFT	Total
7m	minutes	minutes	smoke tests only
15m	20-40 min	~10 min	~30-50 min
75m	3-7 h	30-60 min	4-8 h
100m	TBD (est. 4-8 h)	TBD	TBD
250m	18-25 h	3-5 h	21-30 h
500m	TBD (16 GB+ GPU)	TBD	TBD
750m	TBD (24 GB+ GPU)	TBD	TBD
1b	TBD (24 GB+ GPU)	TBD	TBD

---

Customisation

Settings and Environment

All tunables live in app/config.py. Override via .env at the project root, or by exporting SAVANA_<FIELD> and SAVANA_<GROUP>__<FIELD> variables. Never hardcode paths, hyperparameters, or dataset ids in code; promote them to settings (see ../GUARDRAILS.md, rule 8).

Adding a Size Variant

Edit app/model_service/variants/registry.py:

• Add a variant_<label>() -> ModelConfig function.
• Register it in the _VARIANTS dict.
• Optionally document VRAM and time budgets in docs/04_variants.md.

Pretrain Corpus

Set SAVANA_TRAIN__CORPUS=fineweb_edu (or wikipedia) or pick the corpus in the TUI train wizard. Loaders live in app/training_service/data/.

Chat Behaviour

• Change default sampling in app/config.py (temperature, top-K, loop budget, max new tokens).
• Adjust the chat template in app/training_service/data/chat_template.py.
• Modify the sampler in app/interface_service/generation/sampler.py.

UI Theme

Colours and styles for the Rich TUI live in app/interface_service/views/theme.py. Views under app/interface_service/views/ control layout and copy.

---

Troubleshooting

Common Issues

• cuda: False: your driver does not match the CUDA 12.4 wheels. Edit the URL in [[tool.uv.index]] inside pyproject.toml (e.g. cu121) and rerun uv sync.
• savana not found: run from inside the venv (uv run savana ...) or activate the venv.
• uv sync stalls on the PyTorch wheel: the CUDA index is 3+ GB. First run takes time; subsequent syncs are cached.
• ModelConfig import errors: check python --version; 3.10+ is required.
• Out-of-memory during training: drop to a smaller size, shrink SAVANA_TRAIN__SEQ_LEN, or lower the batch size in app/config.py.

Performance Tips

• Flash Attention 2: uv pip install flash-attn for faster GQA on 75m and above. Practical on Linux with a recent CUDA toolkit; the Windows build path is fragile. Without it, GQA falls back to torch.nn.functional.scaled_dot_product_attention, which is correct but slower.
• Depth extrapolation: sample with /loops above the training value in the chat REPL to trade latency for reasoning quality.
• Resume runs: use the TUI Resume screen or restart the same train command; the checkpoint store picks up the latest step.

---

Development

Code Rules

See ../GUARDRAILS.md. Highlights:

• 200-line soft limit, 250-line hard limit per Python file.
• One responsibility per file, one per function.
• All tunables in app/config.py. No os.environ.get(...) outside config.py.
• Absolute imports only.
• No comments unless the why is non-obvious. One-line docstrings on public entry points.
• Validate at boundaries, trust internal code.

Documentation Layout

• Technical docs under docs/ (01_overview.md, 02_install.md, 03_workflow.md, 04_variants.md).
• Reserve numbering space: tens for sections, ones for sub-topics.

Tests

No automated test suite ships yet. When adding one, mirror the app/ tree under tests/: each public class or entry-point function gets at least one test. Mock only at boundaries (network, disk, GPU); use real objects where cheap.

Tooling

• uv sync to install or update the venv.
• uv run pytest for tests.
• uv run ruff check . and uv run black . for lint and format.
• uv add <pkg> for runtime deps; uv add --dev <pkg> for dev deps. Never call pip directly.

---

License

MIT. See LICENSE.

---

Acknowledgments

The architecture composes ideas from several papers:

• DeepSeek-V2: Multi-Latent Attention.
• DeepSeekMoE: routed plus shared experts.
• Universal Transformers (Dehghani et al., 2018): ACT halting.
• Relaxed Recursive Transformers (Bae et al., 2024): per-loop LoRA.
• Parcae (Prairie et al., 2026): LTI-stable injection parameters.
• Reasoning with Latent Thoughts (Saunshi et al., 2025): looped depth as implicit chain-of-thought.

And to the tooling that keeps a fresh checkout runnable:

• PyTorch for the training and inference stack.
• transformers for the GPT-2 BPE tokeniser.
• Rich for the terminal UI.
• uv for fast, reproducible dependency management.

---

⭐ Star this repo if you found it useful!