Skip to content

Zenith

Stable · v1.0

A from-scratch generative NLP library — decoder-only transformer language models, causal-LM training, and text generation, built on PyTorch tensor primitives.

GitHubPyPI
pip install "zenith-nlp[all]"

Overview

Zenith is a small, correct, readable implementation of the modern decoder-only stack, held to library standards. The architecture is hand-written and modern; PyTorch supplies only autograd, containers, and optimizers. It is the generative counterpart to Polaris — Polaris encodes, Zenith generates — but fully independent: Polaris interop is an optional extra (zenith-nlp[polaris]), never a dependency.

It works: a 10.7M-parameter from-scratch Llama-style decoder, trained in ~10 minutes on a MacBook (MPS), reaches 2.08 bits/char on tiny-shakespeare — matching the well-known nanoGPT baseline.

Problem

The modern decoder-only stack — rotary embeddings, RMSNorm, SwiGLU, KV caching, speculative decoding, LoRA — is scattered across papers and production codebases too large to learn from. Nothing small and readable takes training, inference, and serving seriously as a library at the same time. Zenith exists to fill that gap.

Implemented (v1.0)

  • Decoder-only transformer (DecoderLM) — configurable Llama-style (RoPE, RMSNorm, SwiGLU, weight-tied embeddings) or GPT-2-style (LayerNorm, learned positions, GELU), with an optional fused SDPA attention path.
  • Tokenizers — a dependency-free byte-level tokenizer and a from-scratch, vectorized byte-level BPE; both lossless.
  • Generation (Generator) — greedy, temperature, top-k, nucleus, beam search, repetition penalty, streaming, and a KV cache with rollback; plus greedy-exact speculative decoding: a small draft model proposes tokens, the target verifies — output provably identical to greedy, with up to 3.7× fewer target forward passes.
  • Training (CausalLMTrainer) — warmup/cosine schedule, gradient clipping, AMP, gradient accumulation, LoRA adapters (zenith.peft), torchrun-native DDP, MLflow tracking, and deterministic runs.
  • Instruction tuning (zenith.instruct) — chat template + supervised fine-tuning with response-only loss masking; zenith chat --instruct runs the resulting mini chat model (documented in a model card).
  • int8 quantization — weight-only, ~4× smaller inference weights.
  • Serving & CLI — FastAPI service with SSE streaming, plus zenith generate, chat, console, serve, eval, and Hydra-configured training and sweeps.

Benchmarks

All measured on laptop-scale hardware (Apple Silicon, MPS) and reproducible from scripts in the repo:

  • Headline — 10.7M Llama-style decoder: 2.08 bits/char on tiny-shakespeare (nanoGPT reference: ~2.1).
  • Architecture ablation — same recipe and size, Llama-style vs GPT-2-style: 2.08 vs 2.11 bits/char, with the modern stack converging roughly 2× faster (best validation at epoch 10 vs 17). The honest caveat: at this scale the floor is set by data, not architecture — the modern stack buys convergence speed.
  • Scaling study — 0.6M → 10.7M params, one recipe: each ~3× in parameters buys −0.19, then −0.06, then −0.01 bits/char, flattening into the ~1 MB corpus's data floor.
  • text8 — the same 10.7M model on a 5 MB subset of the standard char-LM corpus: 1.78 bits/char, reported honestly as a subset.

Design decisions

  • Correctness before speed — the fused SDPA path is numerically equivalent to the reference attention; speculative decoding is greedy-exact by construction.
  • RoPE correct under caching — rotary positions are implemented to be consistent between full-sequence training and one-token-at-a-time cached inference.
  • Everything optional is an extra — MLflow, the CLI, serving, and Polaris interop are lazily imported optional dependencies; the core stays lean.
  • Checkpoints are self-describing — model + tokenizer save and load together, LoRA-aware, with backward compatibility for older architecture configs.

Lessons learned

  • The scaling study's flattening curve was the clearest lesson: past ~5M parameters on a ~1 MB corpus, capacity stops paying — the next unit of quality costs data, not parameters.
  • Exactness constraints make optimizations shippable: "speculative output equals greedy output" and "SDPA equals reference attention" turn performance work into testable engineering.
  • The modern architecture bundle (RoPE, RMSNorm, SwiGLU) bought convergence speed more than final quality at this scale — an honest, measured result that reading papers alone would not have produced.

Roadmap

QLoRA and FSDP for larger-scale training, sweep-result aggregation, and a larger CUDA-scale benchmark beyond the laptop-scale text8 subset.