BHAASM Transformer

From-scratch decoder-only transformer with RoPE, GQA, SwiGLU, and optional MoE.

BHAASM is a clean, educational, production-grade transformer language model built entirely from scratch. No hidden abstractions — every component (attention, positional encoding, feed-forward, routing) is implemented explicitly and is easy to modify.

Features

• RoPE — Rotary Position Embeddings with precomputed frequency tables
• GQA — Grouped Query Attention with optional QK LayerNorm and FlashAttention
• SwiGLU — Gated feed-forward with 8/3 hidden dimension ratio
• MoE — 8 experts, top-2 routing, shared expert, load balance auxiliary loss
• RMSNorm — Pre-attention and pre-FFN normalization
• KV Cache — Preallocated or dynamic for efficient autoregressive generation
• BPE Tokenizer — HuggingFace Rust backend, 32K vocab
• Adafactor Optimizer — Memory-efficient training without manual LR tuning
• Multiple Sizes — 150K to 507M parameter configurations

Quick Start

pip install -e .

# Train a tiny model (10 steps, CPU)
python examples/train_tiny.py

# Generate text
bhaasm-generate --checkpoint checkpoints/tiny_demo.pt --prompt "Hello"

# Count parameters
bhaasm-param-count --config tiny_test

Install

git clone https://github.com/cosmic/bhaasm-transformer.git
cd bhaasm-transformer
pip install -e .

Model Configs

Config	Params	Layers	Heads	KV Heads	Embed Dim
tiny_test	~150K	2	2	1	64
tiny_5m	~5.2M	4	4	2	256
tiny_15m	~10.7M	6	6	3	384
mid_33m	~33M	10	8	4	512
mid_55m	~55M	17	8	4	512
mid_500m	~507M	32	16	8	1024

Architecture

Input Tokens → Embedding → [RoPE] → TransformerBlock × N
                                       ├─ RMSNorm
                                       ├─ GQA Self-Attention (with KV Cache)
                                       ├─ RMSNorm
                                       └─ SwiGLU FFN (or MoE Layer)
                                          → LM Head → Next Token

Each TransformerBlock applies pre-norm RMSNorm before both attention and FFN (Pre-LN). When MoE is enabled, every other FFN is replaced with an 8-expert top-2 routed layer plus a shared expert.

Project Structure

bhaasm-transformer/
├── src/bhaasm/
│   ├── model/          # Transformer, attention, MoE, RoPE, RMSNorm, SwiGLU, KV cache
│   ├── tokenizer/      # BPE + char tokenizer (HuggingFace tokenizers backend)
│   ├── training/       # Training loop, checkpointing, loss, Adafactor optimizer
│   ├── inference/      # Generation, sampling, inference interface
│   ├── data/           # Dataset loading, dataset builder, intent dataset
│   └── eval/           # Evaluation metrics
├── configs/            # YAML training configs
├── examples/           # Quickstart training script
└── tests/              # Pytest test suite

Testing

pip install -e ".[all]"
python -m pytest tests/

License

MIT