Skip to content

README.md

Latest commit

 

History

History
139 lines (98 loc) · 6.08 KB

README.md

File metadata and controls

139 lines (98 loc) · 6.08 KB
Raw

Pion

Dense transformer language model. No experimental tricks, just every proven technique from recent research applied correctly. One of the Acceleron series alongside Tachyon (MoE) and Lepton (experimental).

  • Tachyon — Mixture of Experts, sparse routing
  • Lepton — Experimental (thought tokens, memory, speculative decoding)
  • Pion — Best-practices dense transformer. The boring stuff that works.

Architecture

Everything here is backed by published results at small scale (SmolLM2, Gemma 2, Qwen3):

  • GQA — Grouped Query Attention with 5:1 Q-to-KV head ratio. Better than MHA at this scale.
  • QK-Norm — RMSNorm on Q and K before attention. Keeps training stable.
  • Logit Soft-Capping — From Gemma 2. cap * tanh(x / cap) on attention scores. Can be disabled with --no_soft_cap to use FlashAttention/SDPA instead (much less memory).
  • SwiGLU — w2(silu(w1(x)) * w3(x)). Industry standard FFN.
  • RoPE — Rotary position embeddings, theta=10000.
  • Tied Embeddings — Input/output share weights. Critical when embeddings are 20%+ of params.
  • Chunked Cross-Entropy — Applies lm_head per chunk instead of materializing the full (N, vocab) logits tensor. Saves a ton of memory.
  • KV-Cache — For fast autoregressive generation.

Model Sizes

Size d_model Layers Heads (Q/KV) d_ff Params
nano 384 6 6/2 1024 29M
micro 512 12 8/2 1408 60M
small 640 16 10/2 1792 103M
base 768 20 12/4 2048 165M
large 1024 24 16/4 2816 322M
xl 1536 28 16/4 4096 771M

Training

Pion-small (103M) was pretrained on RunPod across a couple of GPU rentals:

  • Steps 0-10K on an A40 (48GB). Batch 64, grad_accum 1, --no_soft_cap with gradient checkpointing. ~4.3s/step, ~31K tok/s.
  • Steps 10K-22.5K on an RTX A4000 (16GB). Batch 16, ~1.6s/step.
  • SFT (3000 steps) on OpenHermes 2.5 using the same A4000. Cosine LR 2e-5, prompt masking so loss only trains on assistant responses.

Total pretraining: ~22.5K steps on ~750M tokens from PreprocessedMIXED. WSD learning rate schedule (warmup-stable-decay).

Checkpoints are on HuggingFace: SuperGoatScriptGuy/Pion

Pretraining

pip install torch numpy tokenizers tqdm huggingface_hub

Data uses preprocessed .npy shards (shape (sequences, 2048), token IDs). Grab some shards to test:

from huggingface_hub import snapshot_download

allow_patterns = [f'shard_{i:05d}.npy' for i in range(50)]
snapshot_download(
    repo_id='SuperGoatScriptGuy/PreprocessedMIXED',
    repo_type='dataset',
    local_dir='./data',
    allow_patterns=allow_patterns,
)
# Memory test
python train.py --model_size nano --data_dir ./data --batch_size 32 --memory_test --no_soft_cap

# Quick local run
python train.py --model_size nano --data_dir ./data --batch_size 4 --max_steps 500 --no_soft_cap

# Real training (cloud GPU)
python train.py --model_size small --data_dir ./data --batch_size 64 --grad_accum 1 \
    --lr 3e-4 --min_lr 3e-5 --max_steps 30000 --no_soft_cap

# Resume from checkpoint
python train.py --model_size small --data_dir ./data --resume checkpoints/small/pion_small_10000.pt --no_soft_cap

Fine-tuning (SFT)

pip install datasets

python finetune.py --checkpoint checkpoints/small/pion_small_22500.pt \
    --batch_size 16 --grad_accum 1 --lr 2e-5 --max_steps 3000 --no_soft_cap

This downloads OpenHermes 2.5 (~1.5GB) automatically and trains with prompt masking — the model only learns from the assistant's responses, not the user's questions.

Useful Flags

  • --no_soft_cap — Disable logit soft-capping, use SDPA/FlashAttention instead. Much less memory. Recommended for anything above nano.
  • --no_gradient_checkpointing — Faster but uses way more VRAM.
  • --compile — Enable torch.compile.
  • --eval_interval 500 — How often to run eval.
  • --save_interval 1000 — How often to save checkpoints.

Generation

# Single prompt
python generate.py --checkpoint checkpoints/small/pion_small_final.pt --prompt "The meaning of life is"

# Interactive mode
python generate.py --checkpoint checkpoints/small/pion_small_final.pt

# Adjust sampling
python generate.py --checkpoint checkpoints/small/pion_small_final.pt --temperature 1.0 --top_k 100 --top_p 0.95

Chat

python chat.py  # auto-selects latest checkpoint
python chat.py --checkpoint checkpoints/sft/pion_sft_final.pt

Commands: /temp 0.8, /topk 50, /topp 0.9, /tokens 128, /quit

Files

  • model.py — Architecture. GQA, QK-Norm, soft-capping, SwiGLU, RoPE, tied embeddings, chunked CE.
  • data.py — Streaming shard-based data loader for pretraining.
  • train.py — Pretraining loop. WSD schedule, BF16, gradient checkpointing.
  • finetune.py — SFT training loop. Cosine schedule, prompt masking.
  • sft_data.py — OpenHermes 2.5 data loading with chat template formatting.
  • generate.py — KV-cache generation with top-k/top-p sampling.
  • chat.py — Interactive generation wrapper.

Tokenizer

Uses a BPE tokenizer (50K vocab) trained on The Pile. Grab it from Tokenizers or just use the pile_tokenizer.json included here.

Notes

At 103M params, Pion-small is a proof-of-concept. It can produce coherent English after pretraining but instruction following after SFT is limited — you really need 1B+ params for that. The architecture and training pipeline are solid though, and scale up to xl (771M) without changes.

Soft-capping is cool in theory but materializes the full O(N^2) attention matrix, which OOMs on anything above nano. In practice you want --no_soft_cap to use SDPA, which is fused and O(N) memory. Gradient checkpointing is also basically mandatory for batch sizes above ~8 on consumer GPUs.