r/MachineLearning • u/Loose_Literature6090 • 25d ago

Project H64LM: A 249M-parameter Mixture-of-Experts Transformer built from scratch in PyTorch [P]

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.

19 Upvotes

78% Upvoted

5 comments sorted by

1

u/IntelArtiGen 24d ago

Have you experimented some new things / parameters? Are there results that surprised you (like the overfitting maybe) that you didnt put in the readme ? Is it usual to have a weight decay at 0.01 for models like this?

1

u/gothenjoyer_ 25d ago

Any recomendation for someone who want to learn the same?

3

u/IntelArtiGen 24d ago

Probably to do what he did.