root-sys-n/hybrid-mamba-50m

Hybrid Mamba-Attention Language Model (50M)

A hybrid Mamba-Attention language model with ~50M parameters, designed for local preflight testing and Colab training.

Architecture

• Pattern: "MMMA MMMA MMMA" (9 Mamba + 3 Attention layers)
• d_model=512, ffn_hidden=1280 (auto-tuned for 49M-51M params)
• Mamba-lite: d_state=8, d_inner=768, dt_rank=16, d_conv=4
• Attention: 8 heads, head_dim=64, RoPE, causal
• SwiGLU FFN, RMSNorm, tied embeddings

Quick Start

Preflight (local, CPU)

python preflight.py

Training (Colab)

Upload colab.ipynb to Colab and run the 2 cells.

Local smoke test

python -c "import yaml; from src.model import build_model; c=yaml.safe_load(open('config/smoke.yaml')); m=build_model(c); print(m.num_params())"

Files

• config/colab.yaml - production config
• config/smoke.yaml - tiny config for CPU testing
• src/model.py - model architecture
• src/tokenizer.py - BPE tokenizer
• src/data.py - streaming dataset with packing
• src/train.py - training script
• preflight.py - 9 preflight checks
• colab.ipynb - 2-cell Colab notebook