| --- |
| license: mit |
| tags: |
| - qwen3-moe |
| - safetensors |
| - transformers |
| - tinyqwen |
| - validation |
| - test-suite |
| - scratch-trained |
| --- |
| |
| # TinyStories Qwen3-MoE 1.95M (tinyqwen3moe2m) HF Validation Suite |
|
|
| This repository provides an ultra-lightweight Qwen3 Mixture-of-Experts (MoE) model file in Hugging Face / Safetensors format, trained to 100% convergence on the TinyStories dataset and optimized for inference engine testing, MoE routing verification, and automated CI pipelines. |
|
|
| ### Why this repository exists |
| When developing a custom Mixture-of-Experts (MoE) inference engine or optimizing low-level routing layers, debugging with full-sized models significantly slows down development efficiency. This suite offers a true **1.95M parameter scale Qwen3-MoE model**, allowing developers to validate their top-k gates, token dispatchers, shared expert combinations, and YaRN RoPE scaling parameters step-by-step with maximum efficiency and verifiable natural language outputs. |
|
|
| ### Key Validation Targets |
| This model is specifically designed to expose architectural layout and edge-case calculation bugs in MoE pipelines: |
| * **Explicit Layer Topology (No Dynamic Extension):** By explicitly defining `num_attention_heads=8` and `head_dim=32` in the configuration, a physical 256-dimensional `q_proj` layer is built directly. This eliminates the need for implicit, runtime 2x layer expansion logic and ensures a predictable memory layout during weight loading. |
| * **Q-Norm / K-Norm Structure Verification:** Validates the application of Per-head RMSNorm directly to the Query and Key tensors prior to the core attention dot-product computation. This is a crucial native feature of the Qwen3 architecture to ensure mathematical stability. |
| * **Native YaRN RoPE Scaling Integration:** Incorporates a true YaRN (Yet Another RoPE Extension) configuration (`rope_scaling`) with `factor=4.0` and `original_max_position_embeddings=64`. This validates that the inference engine can accurately compute frequency adjustments across an expanded context window. |
| * **True 8:1 GQA Ratio:** Implements an asymmetric configuration containing exactly 8 Query heads and 1 Key-Value head (`num_attention_heads=8`, `num_key_value_heads=1`). This checks that KV caching structures, stride calculations, parallel splits, and index handling process Grouped-Query Attention topologies properly without memory alignment failures. |
| * **Top-1 Gated Routing & Token Dispatch:** Features a micro-MoE topology with exactly 2 local experts, routing 1 expert per token (`num_experts_per_tok=1`) across 6 deep transformer layers. This allows developers to strictly track and trace the accuracy of router logit probability flows, routing masks, and row-major tensor switching logic without heavy hardware memory tracking. |
| * **Shared Expert Isolation:** Implements an explicit Shared Expert configuration with an independent dimension shape (`shared_expert_intermediate_size=128`). This helps verify whether the model correctly adds the non-gated common network baseline output onto the single gated routing expert output path without accumulation alignment bugs. |
| * **Continuous Packed Sequence Handling:** The training pipeline utilizes a high-density token-packing algorithm, concatenating sample token continuous sequences separated strictly by control tags (`<s>` / `</s>`) into exact 256-block chunks. This is ideal for testing internal state boundaries, position indexing shifts, and sequence-packing continuous-attention masks in custom hardware layouts. |
| * **Layer-wise Projection Bias Verification ($\pm 0.2$ Uniform Range):** Injected with a frozen constant random uniform bias range ($\pm 0.2$) inside `q_proj`, `k_proj`, and `v_proj` architectures. If an inference engine's dynamic lookup fails to map or slightly shifts these attention biases, the numerical discrepancy scales exponentially across the 6 network layers, immediately destroying greedy generation text into garbage within a few tokens. |
|
|
| --- |
|
|
| ## 📂 Repository Structure & File Descriptions |
|
|
| The repository structure follows standard Hugging Face native configurations for standalone tensor loading: |
|
|
|
|
| ``` |
| |
| ```text |
| SUCCESS |
|
|
| ```text |
| . |
| └── hf/ |
| ├── config.json |
| ├── generation_config.json |
| ├── model.safetensors |
| ├── tokenizer_config.json |
| ├── special_tokens_map.json |
| └── tokenizer.json |
| |
| ``` |
|
|
| ### Hugging Face Native Format (`./hf/`) |
|
|
| * **`hf/model.safetensors`**: The raw, unquantized model weights containing standard query/key/value projections with biases, routing router layers, and MoE experts stored in secure Safetensors format. |
| * **`hf/config.json`**: The architectural configuration file defining MoE hyperparameters (6 layers, 8 heads, 2 local experts, 1 active expert, shared parameters, weight-tying, standard dimensions, and YaRN scaling parameters). |
| * **`hf/generation_config.json`**: Default parameters optimized for text generation. |
| * **`hf/tokenizer_config.json`**: Tokenizer behavior layout specifying the custom ChatML/Qwen3 fast tokenizer setup. |
| * **`hf/special_tokens_map.json`**: Architectural mappings tying special characters to the token blocks. |
| * **`hf/tokenizer.json`**: The custom Byte-Level BPE tokenization descriptor layout trained with a base size of 1000. |
|
|
| --- |
|
|
| ## 🚀 Usage Examples |
|
|
| ### Loading Hugging Face Formats via Python |
|
|
| To match the dataset structure used during training, encode text with `add_special_tokens=False` and manually prepend the exact BOS token ID (`1000`). This ensures perfect token alignment and accurate MoE forward pass sampling. |
|
|
| ```python |
| import torch |
| from transformers import PreTrainedTokenizerFast, Qwen3MoeForCausalLM |
| |
| repo_id = "shibatch/tinyqwen3moe2m" |
| |
| # Load via PreTrainedTokenizerFast to preserve the vocabulary configuration safely |
| tokenizer = PreTrainedTokenizerFast.from_pretrained(repo_id, subfolder="hf") |
| model = Qwen3MoeForCausalLM.from_pretrained(repo_id, subfolder="hf") |
| |
| prompt = "Once upon" |
| |
| # Tokenize without injecting automatic special tokens |
| input_ids = tokenizer.encode(prompt, add_special_tokens=False) |
| |
| # Manually prepend the exact BOS token ID (1000) to match the training pipeline layout |
| input_ids = [tokenizer.bos_token_id] + input_ids |
| inputs = {"input_ids": torch.tensor([input_ids])} |
| |
| with torch.no_grad(): |
| outputs = model.generate( |
| **inputs, |
| max_new_tokens=100, |
| do_sample=False, # Forces deterministic greedy decoding |
| repetition_penalty=1.0, |
| top_p=1.0, |
| bos_token_id=tokenizer.bos_token_id, |
| eos_token_id=tokenizer.eos_token_id, |
| pad_token_id=tokenizer.pad_token_id |
| ) |
| |
| print(tokenizer.decode(outputs[0], skip_special_tokens=True)) |
| |
| ``` |
|
|
| --- |
|
|
| ## 📝 Model Specifications |
|
|
| The routing configuration, hidden layers, and attention mechanics conform strictly to the standard Qwen3-MoE structural profiling. |
|
|
| * **Architecture:** Qwen3-MoE (`Qwen3MoeForCausalLM`) |
| * **Dataset:** TinyStories |
| * **Total Parameters:** 1,954,816 parameters (~1.95M) |
| * **Vocabulary Size:** 1,024 (Custom Byte-Level BPE Tokenizer with 1000 base tokens + special characters) |
| * **Hidden Size (`hidden_size`):** 128 |
| * **Head Dimension (`head_dim`):** 32 (8 heads $\times$ 32 dim = 256, explicitly defining the 256-dimensional `q_proj` from the start without dynamic runtime extensions) |
| * **Number of Hidden Layers (`num_hidden_layers`):** 6 |
| * **Number of Attention Heads (`num_attention_heads`):** 8 |
| * **Number of Key-Value Heads (`num_key_value_heads`):** 1 (Standard GQA 8:1 topology) |
| * **Intermediate Size (`intermediate_size`):** 256 |
| * **Max Position Embeddings (`max_position_embeddings`):** 256 |
| * **Attention Bias (`attention_bias`):** True (Explicitly uniform random between -0.2 and 0.2 for q_proj, k_proj, and v_proj) |
| * **Total Local Experts (`num_experts`):** 2 |
| * **Experts Selected per Token (`num_experts_per_tok`):** 1 (Top-1 Routing) |
| * **Expert Intermediate Size (`moe_intermediate_size`):** 256 |
| * **Shared Expert Intermediate Size (`shared_expert_intermediate_size`):** 128 |
| * **RMS Norm Epsilon:** 1e-06 |
| * **RoPE Base Frequency (`rope_theta`):** 1,000,000.0 |
| * **RoPE Scaling (`rope_scaling`):** `{"type": "yarn", "factor": 4.0, "original_max_position_embeddings": 64}` |
| * **Weight Tying (`tie_word_embeddings`):** True |
|
|
| ## 📜 Acknowledgments & License |
|
|
| * **Original Architecture:** Qwen3 Model Family. |
| * **Dataset:** TinyStories dataset. |
| * **License:** **MIT License**. You are free to use, modify, and distribute these assets for any purpose, commercial or private. |
|
|
