This repository demonstrates a small Mixture-of-Experts language model trained from scratch for educational and research purposes.

Stories-SLM 2 - MoE🤖

This model is a part of a collection of Small Language Models pretrained from scratch on the Tiny Stories Dataset. The collection contains 3 pretrained models (at this moment), more on the way. The model variants in the collection ranges from standard GPT to Mixture-Of-Experts versions built with RoPE, Group Query Attention, and RMSNormalization.

Stories-SLM 2-MoE model improves validation loss while activating only 2 experts per token.

Model Name: Stories-SLM 2 - MoE

Model Description

Stories-SLM 2 - MoE is an advanced small language model pretrained from scratch on the Tiny Stories Dataset. It has 127 million parameters and is trained for 5000 steps on a single Tesla T4 GPU. It is trained on the next token prediction task using Cross-Entropy Loss over 457M tokens. It follows a sparse Mixture-Of-Experts architecture with top 2 of 4 experts being activated for each token. The experts are SwiGLU stype 3-linear layers. Each functionality involving router design, top-k gating and expert dispatch per token is coded FROM SCRATCH.

• Developed by: Namrata Thakur
• Model type: Text Generation
• Language(s) (NLP): English
• License: MIT
• Training Type: Pretraining

Model Sources

• Repository: GitHub Repo
• Demo [optional]: [More Information Needed]

How to Get Started with the Model

AutoModelForCausalLM compatibility coming soon, till then please use below steps to use the model

To install Stories-SLM 2 - MoE, follow these steps:

# Clone the repository
git clone https://github.com/NamrataThakur/Large_Language_Model_From_Scratch_Implementation.git

#Create an environment:
python -m venv env

# Install the required packages
pip install -r requirements.txt

Uses

Stories-SLM 2 - MoE can be used to generate small, grammatically and semantically coherent, creative and complex short stories covering topics like adventure + danger, dialogue, moral stories, object interaction, fear → friendship among other themes.

Chainlit Interface 🖥️

The easiest way to interact with Stories-SLM 2 - MoE is through its Chainlit interface:

chainlit run app_pretrain.py

This will launch a web application where you can input text and see the model's generated responses. A choice can be made on the web interface to choose between Stories-SLM , Stories-SLM 2 and Stories-SLM 2 - MoE models.

Downloading from Huggingface 🤗

To interact with the Stories-SLM 2 - MoE by downloading from the huggingface:

• Step 1: Clone the repo in the local
• Step 2: Below mentioned

from transformer_blocks.gpt2_moe import MoEGPT2
from gpt_Pretraining.text_generation import Text_Generation
import torch

model = MoEGPT2.from_pretrained("NamrataThakur/Small_Language_Model_MOE_127M_Pretrained")
model.eval()

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

#Specifically disable the router noise for inference:
for module in customGPT_pretrain_moe.modules():
  if module.__class__.__name__ == "Router":
      module.noise = False

#---------- Checking the generation to make everything is okay ---------------------------
generation = Text_Generation(model=model, device=device, tokenizer_model='gpt2', 
                                          arch_type='MOE')
start_context = "One day, a "
response = generation.text_generation(input_text=start_context, max_new_tokens = 160, temp = 0.5, top_k=10, kv_cache=False)
print(response)

Model Architecture and Objective

Stories-SLM 2 - MoE uses a standard GPT decoder-only transformer architecture with:

• Attention Type: Group Query Attention
• Num KV Groups: 4
• Normalization: RMSNormalization
• Position Embeddings: Rotary Positional Embeddings (RoPE)
• Num transformer blocks: 8
• Num attention heads: 12
• Embedding dimensions: 768
• Vocabulary size: 50,257 tokens
• Context window: 256 tokens
• Feed-Forward Hidden Dimension: 1024
• Num experts: 4
• Num active experts: 2
• Expert Style: SwiGLU-style 3-linear expert
• Parameters: ~127M (126.75M exact)
• Attention Dropout: 0.2
• Feed-Forward Dropout: 0.2
• Token Dropout: 0.03
• Weight tying between token embeddings and output head

Optimization Config:

• Optimizer: AdamW
• Weight Decay: 0.1
• Beta1: 0.9
• Beta2: 0.95
• Warmup Steps: 330 steps
• Total Steps: ~5000 steps
• use_gradient_clip: True
• Initial Learning Rate: 0.00003
• Maximum Learning Rate: 0.0003
• Gradient Accumulation Steps: 16
• Batch Size: 16
• Global Batch Size: 256
• Scheduler: Linear Increase, followed by Cosine Annealing

Training Details

Training Data

The model was trained on the TinyStories dataset, a collection of short stories designed for training language models. This dataset provides simple narratives that help the model learn coherent story generation while maintaining a smaller size compared to larger language models.

Training Procedure

Stories-SLM 2 - MoE was trained using PyTorch on the TinyStories dataset. The training process involved:

1. Tokenizing the input text
2. Creating sliding windows of fixed block size
3. Training the model with cross-entropy loss
4. Applying learning rate scheduling with warmup and cosine decay

Training Plots

• Learning Rate Vs Steps:

• Loss Vs Steps:

Inference

During inference, Stories-SLM 2 - MoE uses several techniques to produce high-quality text:

• Temperature scaling for controlling randomness
• Top-k sampling for focus and diversity
• Efficient token generation one at a time
• Max New Tokens to determine generation length
• KV Cache for efficient autoregressive generation

Results

Environmental Impact

Carbon emissions can be estimated using the Machine Learning Impact calculator presented in Lacoste et al. (2019).

• Hardware Type: Single Tesla-T4 16GB
• Hours used: [More Information Needed]
• Cloud Provider: Lightning-AI

License

This project is licensed under the MIT License - see the LICENSE file for details.

Support ❤️

If you find Stories-SLM useful, please consider starring the repository ⭐