SLM -- 323M Parameter Small Language Model

A high-performance decoder-only transformer built from scratch for single-GPU training.

Architecture

• 323M parameters (24 layers, 1024 hidden, 2730 FFN)
• Hybrid Attention: alternating Sliding Window (2048) and Global layers
• GQA: 16 query heads / 8 KV heads (2:1 ratio)
• SwiGLU activation, RoPE positional encoding (theta=500k)
• QK-Norm for attention logit stability
• Multi-Token Prediction (DeepSeek V3 style)
• Custom Triton Kernel for ring-buffer KV cache attention

Project Structure

src/
  config.py          -- Model hyperparameters
  model.py           -- SLM top-level class + MTP module
  attention.py       -- HybridAttention (5-path dispatch)
  blocks.py          -- SLMBlock (transformer layer)
  kv_cache.py        -- Zero-copy circular ring-buffer KV cache
  rope.py            -- RoPE frequency precomputation
  norm.py            -- RMSNorm
  activation.py      -- SwiGLU
  muon.py            -- MuonClip optimizer (Muon + AdamW + QK-Clip)
  dataset.py         -- Binary memmap dataset
  tokenizer.py       -- SentencePiece BPE wrapper
  trainer.py         -- Training loop with AMP, grad accumulation, wandb
  kernels/
    __init__.py      -- Triton availability guard
    ring_attention.py -- Custom FlashAttention-2 kernel for ring-buffer caches

scripts/
  download_data.py   -- Download raw parquet data from HuggingFace
  train_tokenizer.py -- Train SentencePiece BPE tokenizer (32k vocab)
  prepare_bin_data.py -- Tokenize raw text into binary format
  train_slm.py       -- Training entry point

tests_modular/
  test_ring_attention.py -- Triton kernel correctness tests (45 cases)

Quick Start

# 1. Install dependencies
uv sync

# 2. Download data
uv run scripts/download_data.py

# 3. Train tokenizer
uv run scripts/train_tokenizer.py

# 4. Prepare binary data
uv run scripts/prepare_bin_data.py

# 5. Train
uv run scripts/train_slm.py --use_wandb --muon_lr 0.02 --adam_lr 3e-4

Training Details

See TRAINING.md for configuration options and hardware recommendations.

Architecture Details

See ARCHITECTURE_REPORT.md for the full technical breakdown.