Project Overview (README)







Architecture · Headline metrics · Quick start · Documentation · Results · References

Status: Architecture, training pipeline, and inference paths are implemented and smoke-tested; the Chinchilla-optimal 8.0B-token pretraining run has not yet started.

Conceptual notes extracted from the source tree live in docs/; the authoritative attention deep-dive is docs/concepts/attention-sinks.md. Every code-symbol citation (file.py:Symbol) is machine-verified by tests/test_doc_refs.py.

Overview#

GPT-OSS-Lite is a from-scratch PyTorch reimplementation of OpenAI's GPT-OSS model (Apache 2.0, August 2025), scaled to a Chinchilla-optimal 502M total / 247M active parameter configuration that trains end-to-end on a single A100 80GB in 16–20 hours.

It is the first long-context MoE and the first attention-sink project in the CoreProjects LLM family, filling two empty cells in the attention-mechanism matrix of frontier-from-scratch reproductions.

Why does this exist? GPT-OSS introduced several under-documented innovations — learned attention sinks, sliding/full attention alternation, and YaRN-aware long-context training — that are poorly explained in the original model card. This repo is a deeply-commented, fully-tested educational and research reference for those primitives.

How it compares to the rest of the portfolio#

Headline metrics#

Both metrics are measured, not assumed. Reproduce with scripts/kv_cache_benchmark.py and scripts/passkey_eval.py.

Why these metrics matter. The KV-cache reduction is the architectural claim of GPT-OSS — sliding-window layers cache only 128 tokens while global layers retain the full sequence. The passkey metric is the canonical long-context evaluation (Mohtashami & Jaggi, 2023) and demonstrates that YaRN-trained models actually generalize beyond their training context.

🗺️ Visual Architecture Atlas#

Explore the full Interactive Visual Systems Guide: four verified Archify showcase maps, live KV-cache memory calculator, MoE router inspector, and verification receipts.

Figure 1: GPT-OSS-Lite Architecture Map — 12-layer decoder alternating between 128-token Sliding Window Attention and Full Attention, YaRN 128K context scaling, and top-2 of 8 MoE. Click image to open interactive guide.

Interactive Architecture & Systems Diagrams#

Architecture#

A 12-layer decoder-only transformer. Every layer alternates between two attention patterns:

Input tokens (vocab = 128,000)
    |
    v
Embedding (d_model=768)              <- weight-tied with output head
    |
    v
12 x GPT-OSS Blocks (gradient checkpointing every 3rd):
    +------------------------------------------------------------+
    |  RMSNorm -> Attention (alternating SWA/full + sink + YaRN)  |
    |  -> Residual -> RMSNorm -> MoE (top-2 of 8) -> Residual     |
    +------------------------------------------------------------+
    |
    v
Final RMSNorm -> Linear head -> Chunked Cross-Entropy (chunk=4096)

Per-layer components#

Configuration#

The canonical config is configs/pretrain_a100_502m.yaml:

Model#

Training#

Quick start#

1. Install#

git clone https://github.com/atandra2000/GPT-OSS-Lite.git
cd GPT-OSS-Lite
pip install -r requirements.txt

2. Verify the architecture (CPU-friendly)#

python3 -m pytest tests/ -v
# 207 tests: 152 passed / 55 skipped across 13 files
# 53 skips = data pipeline (needs sibling LLM/shared_data repo);
# 2 skips = GPU-gated Triton parity (needs CUDA + triton)
# Includes: sliding-window correctness, sink bias, YaRN extrapolation,
# MoE routing, aux loss, gradient flow, checkpoint round-trip, NaN guard

3. Reproduce the headline metric#

python3 scripts/kv_cache_benchmark.py
# HEADLINE METRIC PASSED: 1.94x-2.0x KV-cache reduction

4. Benchmark on GPU#

python3 scripts/microbench_a100.py
python3 scripts/step_time_a100.py --steps 20 --warmup 5

5. Launch a full pretraining run#

python3 training/pretrain.py \
    --config configs/pretrain_a100_502m.yaml \
    --seed 42

6. Resume from checkpoint#

python3 training/pretrain.py \
    --config configs/pretrain_a100_502m.yaml \
    --seed 42 \
    --resume-from 40000   # restores weights + optimizer + scheduler + RNG

Documentation#

Hosted docs site — auto-deployed on every push to main.

Full technical references live in docs/: consolidated concept chapters (theory + implementation), a config/API reference, operation guides, and dedicated training and inference chapters. Every code symbol is cited as a machine-verified file.py:Symbol anchor.

Concepts#

References, guides, training, inference#

Validate docs: python3 scripts/check_docs.py (links) + python3 tests/test_doc_refs.py --strict-coverage (symbol alignment)

Results#

KV-cache reduction (BF16, head_dim=96, batch=1)#

Values are the exact output of scripts/kv_cache_benchmark.py (window=128, batch=1, BF16). The reduction is ≈1.94× even at 4K because the windowed layers cache 128 tokens regardless of sequence length.

Passkey retrieval at 128K (4K-trained model)#

Results pending the first full 8B-token run.

Design decisions#

Project structure#

GPT-OSS-Lite/
├── .github/workflows/
│   └── deploy-docs.yml                 # GitHub Pages auto-deploy
├── configs/
│   └── pretrain_a100_502m.yaml        # canonical training config
├── models/
│   ├── rotary.py                       # RoPE helpers (apply_rope, prune)
│   ├── yarn.py                         # YaRN RoPE scaling
│   ├── attention.py                    # SWA + full + learned sink bias
│   ├── moe.py                          # top-2 routed + 1 shared + aux loss
│   ├── moe_triton.py                   # opt-in fused W1/W3+silu Triton dispatch
│   └── transformer.py                  # top-level GPTOSS + ModelConfig
├── training/
│   └── pretrain.py                     # full training loop + resume
├── inference/
│   ├── generate.py                     # mixed KV-cache generation
│   └── long_context.py                 # 128K passkey retrieval evaluator
├── utils/
│   ├── checkpoint.py                   # atomic safetensors
│   ├── logging.py                      # WandB-capable training logger
│   └── memory.py                       # VRAM estimator
├── data/
│   ├── prepare_data.py                 # shim → LLM/shared_data universal pipeline
│   └── shared_data/                    # vendored universal 8.0B-token pipeline
├── assets/
│   ├── style.css                       # documentation design system
│   └── portal.js                       # portal interactivity
├── scripts/
│   ├── build_docs_html.py              # Markdown → HTML renderer (stdlib only)
│   ├── kv_cache_benchmark.py           # headline metric
│   ├── passkey_eval.py                 # headline metric
│   ├── microbench_a100.py
│   ├── step_time_a100.py
│   ├── e2e_gpu_smoke.py
│   └── check_docs.py
├── tests/                              # 207 tests: 152 passed / 55 skipped, 13 files
│   ├── test_attention.py
│   ├── test_yarn.py
│   ├── test_moe.py
│   ├── test_moe_triton.py
│   ├── test_models.py
│   ├── test_smoke.py
│   ├── test_training.py
│   ├── test_inference.py
│   ├── test_utils.py
│   ├── test_data_pipeline.py
│   └── test_validation.py
├── docs/                               # canonical docs — see docs/README.md
│   ├── README.md                       # doc index + learning path + size table
│   ├── concepts/                       # consolidated theory + architecture
│   ├── references/                     # config + API reference
│   ├── guides/                         # getting-started, operations
│   ├── training.md                     # training loop + data pipeline + YAML reference
│   └── inference.md                    # generation + long-context eval
├── AGENTS.md
├── SKILLS.md
├── LICENSE                             # Apache 2.0
├── requirements.txt
└── pytest.ini

Reproducibility#

Full bit-exact training reproducibility is supported:

• --seed N seeds torch, torch.cuda, numpy, and Python's random.
• Checkpoint RNG state is stored alongside weights (rng_step_N.pt) and restored on resume.
• Deterministic MoE dispatch via torch.argsort(stable=True).
• CUBLAS_WORKSPACE_CONFIG=:4096:8 is set automatically.
• Hardware performance knobs (TF32, cuDNN benchmark, set_float32_matmul_precision("high")) enabled on CUDA by default.
• torch.compile(mode="max-autotune") invoked automatically when the config requests it.

Verification#

# Full test suite (CPU-friendly, ~50 s)
python3 -m pytest tests/ -v
# 207 tests: 152 passed / 55 skipped across 13 files
# With the sibling LLM/shared_data repo checked out: 205 passed / 2 skipped

# Doc-code alignment: every `file.py:Symbol` anchor resolves, every public
# symbol in models/ + training/ + inference/ + utils/ is anchored
python3 tests/test_doc_refs.py --strict-coverage
# Coverage: 100%

# Doc link/stale-pattern lint
python3 scripts/check_docs.py
# check_docs: OK (N files)

# Build the HTML docs portal locally
python3 scripts/build_docs_html.py
# Output: docs_html/

# Headline benchmark
python3 scripts/kv_cache_benchmark.py
# HEADLINE METRIC PASSED: 2.00x KV-cache reduction

Contributing#

PRs welcome for:

• New attention primitives within the GPT-OSS family (e.g., grouped sliding windows, hierarchical sinks)
• Aux-loss variants (router-z loss, expert capacity factors)
• Long-context benchmarks (RULER, LongBench, needle-in-a-haystack variants)
• Tokenizer swaps with documented re-derivation of yarn_target_seq_len

Please:

1. Read docs/concepts/attention-sinks.md before touching models/attention.py.
2. Run pytest tests/ -v — all tests must pass (currently 152 passed / 55 skipped; 205 passed / 2 skipped with the sibling LLM/shared_data repo present).
3. If you touch docs or rename symbols, run python3 tests/test_doc_refs.py --strict-coverage and python3 scripts/check_docs.py — stale anchors fail.
4. Run scripts/kv_cache_benchmark.py and confirm the 2.0× reduction still holds.
5. Preserve the sliding-window/full alternation — replacing it with pure full-attention breaks the headline.

Known caveats#

• Full 8B-token pretraining run not yet started (no GPU on dev machine). The 207-test suite (152 passed, 55 skipped without the sibling LLM/shared_data repo; 205 passed, 2 skipped with it) validates all primitives on CPU + tiny shapes.
• The 53 data-pipeline tests need the sibling LLM/shared_data repo. Without it they collect and skip individually — they are not silently dropped, so a bare checkout reports 55 skips rather than 2.
• passkey_eval.py requires a trained checkpoint; it runs as a stub on untrained models.
• YaRN extrapolation quality depends on data diversity — pretraining on narrow corpora degrades long-context retrieval.

References#

• GPT-OSS model card — OpenAI, August 2025
• Raschka, "From GPT-2 to GPT-OSS: Analyzing the Architectural Leap" — Sep 2025
• StreamingLLM (attention sinks) — Xiao et al., arXiv:2309.17453
• Off-by-one attention — arXiv:2402.09093
• YaRN — Peng et al., arXiv:2309.00071
• Longformer (sliding window) — Beltagy et al., arXiv:2004.05150
• DeepSeekMoE — Dai et al., arXiv:2401.06066
• Chinchilla scaling laws — Hoffmann et al., arXiv:2203.15556
• Passkey retrieval benchmark — Mohtashami & Jaggi, 2023

License#

Apache 2.0 — matches the GPT-OSS upstream license. See LICENSE.