iq1xxs/MoE-MicroGPT-100M

Mamba-MoE-100M (Conversational Edge Prototype)

This repository contains the weights and architecture code for an experimental, ultra-lightweight language model based on the State Space Model (SSM) Mamba architecture combined with Mixture of Experts (MoE) routing topology.

The architecture is designed specifically as an experiment for computational efficiency on edge devices, low-power embedded systems, and resource-constrained environments.

📊 Architectural Specifications

• Base Architecture: State Space Model (Mamba-1 / Selective Scan) providing linear context complexity.
• MoE Topology: Token-level balancing routing mechanism managing 8 local experts.
• Expert Activation: Dynamic Top-3 routing strategy (strictly 3 active experts out of 8 per token).
• Hidden Size ((d_{model})): 512
• Vocabulary Size: 50,260 tokens (Optimized English Byte-Pair Encoding).
• Model Weight Size: ~231 MB (Optimized float16 state_dict format).
• Training State: The checkpoint has completed its initial calibration phase on conversational instruction data (no_robots dataset) reaching a Cross-Entropy Loss threshold of ~5.5.

📦 Repository Structure

• potatomoe-100m-english.pt — Pre-trained float16 PyTorch state_dict checkpoint.
• moe_mamba_tiny.py — Core architecture implementation including SSM layers, MoE routing blocks, and Auxiliary Load Balancing Loss function.
• tokenizer.json — Trained BPE tokenizer configuration with embedded control tokens (<|system|>, <|user|>, <|assistant|>).

💻 Baseline Inference (PyTorch)

To execute the model framework, initialize the unified container structure with the weights provided in this repository.

⚠️ CRITICAL NOTE ON FILENAMES: Rename the downloaded weights file from moe-microgpt-100m.pt to potatomoe-100m-english.pt if required for internal resolution, and ensure proper module aliasing as detailed in the repository configuration.

Full implementation details and code can be found in the referenced repository.

To execute the model framework, initialize the unified container structure with the weights provided in this repository. Note that moe_mamba_tiny.py should be imported or aliased locally as modeling_potatomoe.py for proper execution.

import torch
from transformers import PreTrainedTokenizerFast
from modeling_potatomoe import MambaTop3MoELayer

class UnifiedMambaMoE(torch.nn.Module):
    def __init__(self, d_model=512, num_experts=8, num_active_experts=3, vocab_size=50260):
        super().__init__()
        self.embedding = torch.nn.Embedding(vocab_size, d_model)
        self.mamba_moe = MambaTop3MoELayer(d_model=d_model, num_experts=num_experts, num_active_experts=num_active_experts)
        self.lm_head = torch.nn.Linear(d_model, vocab_size, bias=False)

    def forward(self, input_ids):
        x = self.embedding(input_ids)
        x, importance = self.mamba_moe(x)
        logits = self.lm_head(x)
        return logits, importance

🚀 Quick Start (Inference)

1. Install dependencies:

   pip install torch transformers
   

2. Run the model: Create a run.py file in the same folder with your downloaded model files and weights, then execute:

   python run.py
   

---

⚠️ Note on File Names (If you rename files)

The inference script expects the files to match the names in this repository. If you have renamed your files, please make sure they match inside your run.py:

• Weights file (.pt): If your weights are named potatomoe-100m-english.pt instead of moe-microgpt-100m.pt, look at line 23 (or where weights = ... is defined) and change it to your filename.
• Architecture file (.py): If you renamed moe_mamba_tiny.py to modeling_potatomoe.py, make sure your import at the top of run.py reflects this change:

  from modeling_potatomoe import MambaTop3MoELayer
  

📜 License

This project is released under the Apache 2.0 License.

---