
H64LM: A 249M-parameter Mixture-of-Experts Transformer built from scratch in PyTorch
Hi everyone,
I built H64LM, a research project to better understand modern LLMs by implementing one from scratch in PyTorch.
Instead of relying on high-level training frameworks, I implemented the core components myself attention, MoE routing, normalization, and the training loop.
Features
- 249M-parameter Transformer
- Grouped Query Attention (GQA)
- Sparse Mixture-of-Experts (8 experts, Top-2 routing) with 3 auxiliary routing losses
- SwiGLU, RoPE, RMSNorm
- Sliding-window attention
- Mixed-precision training, gradient accumulation
- Custom training loop (no Trainer abstractions)
- Checkpointing and resume support
The included checkpoint was trained on a subset of WikiText-103 to validate the pipeline end-to-end, not to be a strong model it's visibly overfit past epoch 10 (best val PPL ~40.5).
Known limitations are documented in the README, including batch-size-1-only generation and no true DDP (falls back to DataParallel).
GitHub: https://github.com/Haiderkhan64/H64LM
Feedback on the implementation or architecture is very welcome.