Tiny Qwen MoE – TinyStories
A small, decoder-only, Mixture-of-Experts (MoE) language model implemented from scratch in PyTorch, inspired by the architectural concepts used in Qwen3.
The model was trained on a subset of the TinyStories dataset to explore causal language modelling, sparse expert routing, and modern Transformer components under limited compute resources.
This is an independent implementation, not an official Qwen model or a fine-tuned Qwen checkpoint.
Model Details
| Property | Value |
|---|---|
| Architecture | Decoder-only Transformer |
| Model type | Sparse Mixture-of-Experts |
| Framework | PyTorch |
| Training objective | Causal Language Modelling |
| Tokenizer | Qwen2 tokenizer |
| Training dataset | TinyStories |
| Training approach | From scratch |
Architecture
The model incorporates the following components:
- Grouped Query Attention (GQA): Uses fewer key-value heads than query heads.
- Rotary Positional Embeddings (RoPE): Encodes positional information.
- RMSNorm: Pre-normalization for Transformer blocks.
- SwiGLU: Feed-forward network activation.
- Mixture of Experts: Routes each token to a subset of available experts.
Model Configuration
| Parameter | Value |
|---|---|
| Hidden dimension | 384 |
| Transformer layers | 6 |
| Attention heads | 8 |
| KV heads | 4 |
| Attention head dimension | 48 |
| Expert count | 16 |
| Active experts per token | 2 |
| Expert hidden dimension | 1536 |
| Maximum sequence length | 256 |
Training
The model was trained from randomly initialized weights using the TinyStories dataset.
- Dataset:
roneneldan/TinyStories - Training data: 10,000 examples
- Validation data: 1,000 examples
- Objective: Next-token prediction
- Optimizer: AdamW
- Training environment: Kaggle GPU
The dataset subsets and training configuration were selected to make experimentation feasible with limited compute resources.
Training Results
The model was trained for 5 epochs on a subset of TinyStories.
Training Configuration
- Dataset:
roneneldan/TinyStories - Training examples: 10,000
- Validation examples: 1,000
- Batch size: 4
- Learning rate: 1e-4
- Optimizer: AdamW
- Weight decay: 0.1
- Training epochs: 5
- Maximum sequence length: 256
- Hardware: Kaggle GPU (exact accelerator to be specified)
Loss Curve
Final training loss: To be filled from the training logs.
Final validation loss: To be filled from the training logs.
Inference
You can run inference using the inference.py script provided in the GitHub repository.
Installation
pip install torch transformers huggingface_hub
Generate text
python inference.py \
--model moe \
--repo SA0901/tiny-qwen-moe-tinystories \
--prompt "Once upon a time, there was a little rabbit" \
--max-new-tokens 100
The script automatically downloads the model weights, configuration, and tokenizer from Hugging Face.
Decoding: Greedy autoregressive generation.
Note: This is a base causal language model, not an instruction-tuned conversational model.
Limitations
- Trained on a relatively small subset of TinyStories.
- Limited model capacity and context length.
- English-focused training data.
- Generation quality and generalization are limited by the training scale.
- Not benchmarked against production-scale language models.
Repository Contents
model.py– Model architecture and Transformer components.configs.py– Model configuration.pytorch_model.bin– Trained model weights.- Tokenizer files – Tokenization configuration and vocabulary.
model_config.json– Exported model configuration.
Acknowledgements
- Qwen for architectural inspiration and tokenizer resources.
- TinyStories for the training dataset.
Disclaimer
This project is an independent educational implementation. It is not affiliated with or endorsed by Alibaba Cloud or the Qwen team.