GPT-OSS-Lite

OpenAI GPT-OSS-Lite

From-scratch PyTorch reproduction of OpenAI’s GPT-OSS at ~502M params total / ~247M active — 6 sliding-window (W=128) + 6 full-attention layers in alternation, learned per-head attention-sink bias, YaRN RoPE 128K, and top-2-of-8 MoE with one sanctioned Triton grouped-GEMM. Read it like a field notebook: a name, a wiring sketch, then the measurements.

KV CACHE REDUCTION
≥ 1.8× CUT
6 SWA (W=128 ring buffer) + 6 global layers → O(1) decode per step at 128K
ACTIVE COMPUTE
49.2% FLOPs
Top-2 of 8 routed (2×1536 SwiGLU) + 1 shared ≈ 247M / 502M total params
ATTENTION SINKS
β ∈ [−10, +15] CLAMP
Learned per-head sink bias (virtual sink column, β ∈ [−10, +15]) stabilizes softmax across SWA layers
CONTEXT LENGTH
128K YaRN RoPE
NTK-by-parts ramp θ=100K, scale=32, pruned RoPE on 25% global-layer dims
FIG. A0 ATTENTION-SINK GRAVITY WELL & YaRN 128K RESONANCE SINK · W=128 SLIDING WINDOW · 36 YaRN PHASORS
Ah[q,k] = softmaxk ∈ K∪sink ( q⊤k / √d + βhsink · δk=sink )
Virtual Sink Column (per-head βhsink ∈ [−10, +15], zero K/V) SWA Window (W=128, Ring Buffer) YaRN 128K Extrapolation Manifold CLICK: INJECT QUERY PULSE · HOVER: PROBE PHASOR
01 · params
~502M total
02 · layers
12 alt 6+6
03 · d_model
768
04 · attn
8Q/4KV GQA
05 · ffn
1536 SwiGLU
06 · vocab
128K
07 · ctx
128K YaRN
08 · MoE
top-2/8 +1 shared
GPT-OSS-Lite core mechanisms · interactive benchmark labs
01 · SWA KV CACHE Mixed KV Footprint Analyzer

Compare KV cache memory between full-attention (all 12 layers, full context) and the GPT-OSS-Lite mixed scheme (6 global + 6 SWA W=128 ring buffer) with GQA 4 KV heads × 96 head_dim × bf16.

Context Length: 32,768 tokens
Full Attention (12 layers): 0.56 GB
Mixed (6 Global + 6 SWA W=128): 0.28 GB
KV Cache Saved: 0.28 GB (2.00× cut)
02 · ATTENTION SINKS Per-Head Sink Bias Monitor

Visualize learned per-head additive sink bias βh across 8 GQA heads. Clamped to [−10, +15] at forward time to prevent BF16 SDPA mask-add overflow.

Mean Sink Bias: +0.00
Sink Status: INITIALIZED
03 · MoE DISPATCH Top-2 of 8 Router + Shared

Simulate the GPT-OSS MoE router: top-2 of 8 routed experts (SwiGLU, ffn=1536) + 1 always-active shared expert. Standard aux load-balancing loss (α=0.01).

Active Experts: #02, #05 + shared
Dispatch Backend: torch (default)
FIG. A1 FULL TRAINING STEP PIPELINE FORWARD · AUTOGRAD BACKWARD · ADAMW
• FORWARD · Activations xt → 12 Layers (SWA ↔ Full) → MoE Grouped-GEMM
Forward Activations Backward Gradients (∂ℓ/∂θ) AdamW Update (Δθ) CLICK STAGE: INSPECT · HOVER: TENSOR SHAPES
8 chapters

Part I: Foundations

Ch 1

Project Overview (README)

502M-param pure-PyTorch GPT-OSS reproduction: sliding-window/full attention alternation, learned attention-sink bias, YaRN RoPE 128K, and top-2-of-8 MoE with one sanctioned Triton kernel.

 Ch 2

Foundations & 12-Layer Architecture

The alternating 6 SWA + 6 full-attention stack, GQA, RMSNorm, and the full wiring.

 Ch 3

Attention Sinks & Sliding-Window Stabilization

Why learned per-head sink bias stabilizes SWA — the rolling-buffer ↔ softmax cliff problem and the clamp.

 Ch 4

Attention Geometry & YaRN RoPE

Pruned RoPE on global layers, YaRN NTK-by-parts ramp at θ=100K, and the 128K extrapolation math.

 Ch 5

Sliding-Window / Full Alternation

One-topic primer: which layers slide, the banded mask, and the 2.00× KV-cache cut.

 Ch 6

Learned Attention Sinks

One-topic primer: softmax mass under eviction, the zero-value sink key, and the clamp.

 Ch 7

YaRN Scaling to 128K

One-topic primer: the frequency ramp, mscale temperature, and per-layer RoPE pruning.

 Ch 8

Mixture of Experts — Top-2 of 8 + 1 Shared

Top-2 of 8 routed + 1 shared, standard aux loss, grouped dispatch, expert-load balance.
1 chapter

Part II: The data

Ch 9

Tokenization — Tiktoken o200k

Tiktoken o200k_base, vocab 128000, the shared 8-B token data pipeline contract.
6 chapters

Part III: The training run

Ch 10

Getting Started — From Zero to a Running Loop

From zero to a running training loop — install, verify the math, full run, resume.

 Ch 11

Training, Memory Stack & Data Pipeline

Mixed-precision stack, the gradient checkpointing contract, and the long-context data path end to end.

 Ch 12

Triton MoE Kernels & Memory Checkpointing

Sanctioned fused W1/W3+silu grouped-GEMM, gradient checkpointing per 3rd layer, NaN guard.

 Ch 13

Optimizers & Numerical Stability

BF16 + FP32 AdamW master, TF32, manual FP32 attention, sink-bias clamp, chunked CE.

 Ch 14

Operations Guide — Launch, Monitor, Resume

Launch, monitor, NaN recovery, and resume for a production pre-training run.

 Ch 15

Inference & 128K Long Context

MixedKVCache (ring buffer + exponential growth), 128K passkey retrieval, and O(1) decode per step.