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.

Model	Params	Architecture	Validation Loss
Stories-SLM	53M	Dense - MHA	1.78
Stories-SLM 2	48M	Dense - GQA	1.73
Stories-SLM 2-MoE	127M	Sparse - Mixture-of-Experts	1.67

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:

bash

# 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

bash

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:

Tokenizing the input text
Creating sliding windows of fixed block size
Training the model with cross-entropy loss
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 ⭐