Update README.md

README.md

@@ -15,4 +15,183 @@ tags:
 - llm
 - text-generation
 - causal-lm
----
+---
+# Model Card for Tiny-MoE
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+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
+
+<!-- Provide a longer summary of what this model is. -->
+
+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
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** https://github.com/AbdelrhmanEbied/Tiny-MoE
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+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
+
+<!-- This section is meant to convey both technical and sociotechnical 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
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+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.
+
+```bash
+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
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+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
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the 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
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+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
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+- **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