NanoMoE is a minimal rewrite of Karpathy's NanoGPT, from scratch, with pedogical implementations for some new features, such as MOEs.
A compact, from-scratch character-level Transformer model with Rotary Position Embeddings, Mixture-of-Experts feed-forward layers, and F-gram contextual augmentation — all implemented in a single model.py file. Beats out NanoGPT's generative capabilities on a roughly similar LOC (to offset higher memory but retain gains from MOE and RoPE, consider turning of the F-gram context). Run the (badly named) model.py file to start training; you might want to change some of the hyperparameters to fit on a consumer GPU. Verified trains across CUDA, MPS and ROCM.
Written for practice, manually, on a keyboard in two afternoons.
Currently, it loads the TinyShakespeare dataset; perhaps a switch to FineWeb is warranted.
-
Swap in a subword tokenizer (BPE/WordPiece) in place of the character encoder
-
Automate F-gram mining from the corpus rather than hard-coding
-
Add
muonsupport, add cosine annealing and other LR schedulers. -
Allow mixed-precision training.
-
tritonkernels for fun? -
I'll be adding an arg parser soon, but until then, these are rough, recommended values! (Although I assume TinyShakespeare would work no matter what you choose to run on). Here's a schema that anyone who wants to contribute could follow.
Argument Description Default --embedding_dimToken embedding size 128 --num_headsNumber of attention heads 4 --num_layersNumber of Transformer blocks 4 --block_sizeContext window (sequence length) 64 --dropoutDropout probability 0.1 --moe_expertsNumber of experts in the MoE layer 4 --fgram_max_nMaximum n-gram length for F-gram augmentation 3 --learning_rateAdamW learning rate 3e-4 --batch_sizeBatch size 512 --epochsNumber of training epochs 10