A Qwen-like decoder-only model written from scratch in PyTorch. Not a thin wrap around transformers model classes. I wanted the pieces modern small LMs actually use under my fingers.
Architecture#
Implemented in model.py:
- Grouped Query Attention (GQA): fewer KV heads than query heads, then expand KV for full attention
- RoPE on Q/K for relative position
- RMSNorm instead of LayerNorm
- SwiGLU feed-forward
- causal mask plus an optional KV cache path for generation
Hyperparameters for the Shakespeare run (see config.json) stay small on purpose. This is not full Qwen scale:
| Setting | Value |
|---|---|
n_dim | 512 |
hidden_dim | 768 |
n_heads | 8 |
n_groups | 4 (GQA) |
depth | 12 |
seq_len | 128 |
| tokenizer | GPT-2 vocab (~50k) |
Training and inference#
- Data: Tiny Shakespeare tokenized offline, loaded as a tensor
- Train loop: next-token CE, AdamW, CLI flags for depth/dim/batch (SageMaker-friendly channel paths)
- Inference: load a checkpoint and stream text (temperature, max new tokens)
- Hub: experiment weights at heissanjay/qwen-tiny-shakespeare-experiment
Tokenizer files under tokenizer/ keep encode and decode aligned with training.
Why it exists#
I wanted one stack that owns GQA, RoPE, SwiGLU, and RMSNorm end to end: train, checkpoint, sample, publish. Handy when I want to poke group count, depth, or seq len without first wrestling a huge open-weight codebase.