Model Card for Tiny-MoE

Tiny-MoE is a lightweight Mixture-of-Experts (MoE) decoder-only language model implemented entirely from scratch in native PyTorch, and trained end-to-end on Kaggle's free-tier hardware (2× NVIDIA T4 GPUs).

Model Details

Model Description

Tiny-MoE is an open-source, highly efficient Mixture-of-Experts language model built entirely from scratch using native PyTorch, without relying on high-level modeling frameworks. It was designed to maximize compute efficiency on Kaggle's free-tier cloud hardware and serves as a bottom-up implementation of modern LLM techniques, including Multi-head Latent Attention (MLA), YaRN long-context extension, and optimized MoE routing.

The model is a 14-layer decoder-only Transformer. Each layer combines an MLA attention block with a Mixture-of-Experts feed-forward network made up of 8 routed experts, 1 shared expert, and Top-2 routing. Weight absorption can be enabled at inference time to reduce computation.

• Developed by: Abdelrhman Ebied
• Funded by: Self-funded (trained on Kaggle's free-tier compute)
• Shared by: Abdelrhman Ebied
• Model type: MoE decoder-only Transformer (causal language model)
• Language(s) (NLP): English
• License: Apache License 2.0
• Finetuned from model: Not applicable — trained from scratch (base model)

Model Sources

• Repository: https://github.com/AbdelrhmanEbied/Tiny-MoE

Uses

Direct Use

Tiny-MoE is a base (pretrained) model — it has not been instruction fine-tuned. It can be used directly for raw text generation / completion via the project's inference pipeline, which supports temperature sampling, top-k sampling, top-p (nucleus) sampling, repetition penalty, and n-gram blocking. As a base model, it is not expected to reliably follow instructions or hold a conversation out of the box.

Bias, Risks, and Limitations

Tiny-MoE is primarily a research and learning project, and is currently a base model (no instruction fine-tuning). It should not be expected to follow instructions, hold a conversation, or produce chat-style responses reliably — its outputs are raw text continuations shaped by its pretraining data. The project's primary focus has been implementing and understanding modern LLM architectures and training techniques from scratch, rather than downstream task or chat performance. Beyond this, no formal bias or risk evaluation is documented.

Recommendations

Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. Given the model's small scale, research-oriented training focus, and base (non-instruction-tuned) status, it should not be relied on for production, chat, or safety-critical use cases without further fine-tuning. Future updates are planned to include instruction tuning, preference optimization, and continued training to improve conversational quality.

How to Get Started with the Model

Use the code below to get started with the model.

git clone https://github.com/AbdelrhmanEbied/Tiny-MoE.git
cd Tiny-MoE
pip install -r requirements.txt
!python -m inference.generate

Training Details

Training Data

Tiny-MoE was pretrained on a streaming mixture of three Hugging Face datasets, interleaved with fixed sampling probabilities using the all_exhausted stopping strategy:

Dataset	Purpose	Sampling Probability
FineWeb-Edu (sample-10BT)	General web and educational text	60%
Cosmopedia v2	Synthetic educational corpus	25%
Open-Web-Math	Mathematical reasoning	15%

Training Procedure

Training used a streaming data pipeline that tokenizes incoming text, packs it into a continuous token buffer, and slices fixed-length (512-token) training sequences rather than padding individual documents — maximizing GPU utilization on constrained hardware.

Preprocessing

Text is streamed directly from Hugging Face (no full dataset download), interleaved across the three source datasets, shuffled with a large buffer, sharded across distributed processes and DataLoader workers, tokenized, and packed into contiguous 512-token windows to generate input_ids, position_ids, and shifted labels.

Training Hyperparameters

• Training regime: FP16 mixed precision
• Optimizer: 8-bit AdamW
• Learning Rate: 2.5e-4
• Weight Decay: 0.1
• Betas: (0.9, 0.95)
• Gradient Clipping: 1.0
• Training Steps: 30,000
• Warmup Steps: 2,800
• Micro Batch Size: 64
• Gradient Accumulation: 4
• Effective Sequence Length: 512

Evaluation

Testing Data, Factors & Metrics

Metrics

During training, the following metrics were tracked: training and validation loss, perplexity, router entropy, router confidence, expert utilization, load ratio, load standard deviation, and minimum/maximum expert load, logged via Weights & Biases.

Results

Training Loss: Loss starts around ~10.5 at the beginning of training and drops sharply within the first ~2k–3k steps to roughly 3–4. From there it continues a slower downward trend, settling into a noisy band roughly between 2 and 4 for the remainder of the 30k training steps, ending around ~2.5–3.

Validation Perplexity (val_ppl): Evaluated periodically at checkpoints (labeled Base-1 through Base-11). It starts extremely high (~375 at Base-1, the first checkpoint) and falls steeply over the first several checkpoints — dropping to roughly ~170 (Base-1 end), ~75–115 (Base-2), then under 60 by Base-3, and below 50 from Base-4 onward. From there it continues to decrease gradually and more gently through the remaining checkpoints, trending toward the low-to-mid 40s by the final checkpoints near step 30k. I can't read exact numeric values off the chart with full precision, so treat these as approximate.

Router Confidence: Starts low (~0.2) early in training, consistent with an untrained/undecided router, and rises quickly to roughly 0.4 within the first ~2k–3k steps. It then stays relatively stable in the ~0.35–0.45 range for the rest of training, with periodic dips that align with the same checkpoint boundaries seen in the val_ppl chart — likely corresponding to evaluation or checkpoint-switch points rather than a genuine drop in routing stability.

Router Load Ratio: Fluctuated between roughly 1.2 and 2.0 across training, with typical values sitting around 1.4–1.5 — indicating experts were generally load-balanced with occasional wider imbalance spikes.

Router Entropy: Stayed near the theoretical maximum of ln(8) ≈ 2.08 (the max entropy for 8 routed experts), reflecting near-uniform routing. It occasionally dropped below this but recovered quickly, suggesting temporary routing imbalances rather than persistent collapse.

Expert Utilization: Ranged from about 0.975 to 0.999, typically hovering around 0.99 — indicating experts were consistently and almost fully utilized throughout training.

Training Duration: ~90 hours total

Technical Specifications

Model Architecture and Objective

Tiny-MoE is a 14-layer decoder-only Mixture-of-Experts Transformer trained with a causal language modeling objective. Each layer applies RMSNorm, Multi-head Latent Attention (MLA) with RoPE positional encoding and SDPA attention, followed by RMSNorm and a Top-2-routed MoE feed-forward block (8 routed experts + 1 shared expert) executed via grouped GEMM. Auxiliary router loss and router z-loss are used to stabilize routing during training, and YaRN can be applied to extend context length.

Parameter	Value
Vocabulary Size	32,000
Hidden Size	512
Transformer Layers	14
Attention Heads	8
Routed Experts	8
Shared Experts	1
Experts per Token	2
MoE Intermediate Size	1024
Maximum Context Length	512 (2048 via YaRN)
RMSNorm Epsilon	1e-6
RoPE Theta	10,000
Attention Implementation	SDPA
Weight Tying	Enabled
Total Parameters	~150M–200M
Active Parameters	~70M–90M per token

Compute Infrastructure

The model was developed and trained entirely on Kaggle.

Hardware

• Hardware Type: 2× NVIDIA T4 GPUs
• Cloud Provider: Kaggle

Software

Native PyTorch, DeepSpeed (ZeRO Stage 2, expert parallelism), Hugging Face Datasets, Torch Compile.

Glossary

• MoE (Mixture-of-Experts): An architecture where only a subset of specialized "expert" sub-networks are activated per token, reducing active compute relative to total parameter count.
• MLA (Multi-head Latent Attention): An attention variant that compresses key/value representations into a latent space to reduce memory and compute.
• YaRN: A technique for extending a model's usable context length beyond what it was originally trained on.
• RoPE (Rotary Position Embeddings): A positional encoding method that rotates query/key vectors to encode relative position information.

Model Card Authors

Abdelrhman Ebied