GenakiGyoka/deepseek-v4-nano-tinystories

DeepSeek V4-Nano (TinyStories)

DeepSeek V4-Nano is an efficient 130M-parameter Mixture-of-Experts (MoE) Large Language Model built from scratch using PyTorch. It incorporates modern LLM architectural innovations including Multi-Head Latent Attention (MLA), Auxiliary-Loss-Free Load Balancing (Bias-based routing), and a Shared + Routed Expert Architecture, trained on the TinyStories dataset (~925M tokens).

---

📌 Model Specifications

• Model Architecture: DeepSeek V4-Nano (MoE + MLA)
• Total Parameters: ~130.1 Million
• Active Parameters per Token: ~42.3 Million
• Context Length (seq_len): 512 tokens
• Vocabulary Size: 16,384 tokens (Custom Byte-Pair Encoding BPE)
• Embedding Dimension (dim): 512
• Number of Layers: 8
• Number of Attention Heads: 8
• Total Experts: 8 routed experts + 1 shared expert
• Active Experts per Token (top_k): 2 routed + 1 shared
• Load Balancing: Auxiliary-Loss-Free Bias Adjustment
• Dataset: TinyStories (~925M tokens)
• Training Steps: 19,500 steps (Batch size 32, Grad accumulation 4)
• Best Validation Loss: 0.9236
• License: MIT

---

💡 Key Architectural Features

1. Multi-Head Latent Attention (MLA): Compresses Key-Value cache into low-rank latent vectors, significantly reducing VRAM consumption and inference memory footprint while preserving multi-head attention expressiveness.
2. Shared + Routed MoE: Combines a persistent shared expert that captures general language patterns with 8 specialized routed experts.
3. Auxiliary-Loss-Free Load Balancing: Dynamically updates routing biases during training to maintain uniform expert utilization without requiring auxiliary loss terms.
4. Custom BPE Tokenizer: Trained specifically on TinyStories text to maximize tokenization efficiency for narrative text.

---

🧪 Example Generations (Phase 1 Evaluation)

Below are representative text generations from the model across different temperatures using prompt seeds:

Prompt 1: "Once upon a time,"

Temperature 0.2: "Once upon a time, there was a little girl named Lily. She loved to play outside in the sunshine. One day, she found a shiny rock on the ground. It was so pretty and smooth. She picked it up and held it in her hand. Suddenly, she heard a loud noise. It was her friend, Timmy. He was crying because he lost his toy car. Lily wanted to help him find it. She looked around and saw the toy car under a bush. She picked it up and gave it to Timmy. He was so happy and hugged her. Lily felt happy too."

Temperature 0.5: "Once upon a time, there was a little girl named Lily. She loved to play with her toys and run around in the park. One day, she saw a big dog and she wanted to pet it. But the dog was not friendly and it bit her. Lily's mom took her to the hospital and the doctor said she needed a shot. Lily was scared and didn't want the shot. She cried and cried until her mom gave her a hug and a kiss."

Prompt 2: "Tom walked into the forest."

Temperature 0.2: "Tom walked into the forest. He saw a big tree with a hole in it. Tom looked inside the hole and found a small, shiny box. He opened the box and found a note. The note said, 'To Tom, from the squirrel.' Tom was very happy and thanked the squirrel."

Prompt 3: "The little rabbit said,"

Temperature 0.5: "The little rabbit said, 'I don't like the rain, it makes me wet and cold.' The wise owl said, 'I have an idea. Let's go to my tree and stay dry.' They went to the tree and the wise owl said, 'Let's cover ourselves with my wings so we don't get wet.' The little rabbit was happy and said, 'Thank you, wise owl.'"

---

🚀 Quickstart & Usage

1. Requirements

pip install torch tokenizers

2. Run Inference in Python

import torch
from tokenizers import Tokenizer
from deepseek_v4 import DeepSeekV4

# Load tokenizer
tokenizer = Tokenizer.from_file("tokenizer.json")

# Model configuration
model = DeepSeekV4(
    vocab_size=16384,
    seq_len=512,
    dim=512,
    n_heads=8,
    n_layers=8,
    n_experts=8,
    expert_hidden_dim=1024,
    top_k_experts=2,
    use_shared_expert=True,
    load_balance_method="bias"
)

# Load weights
state_dict = torch.load("pytorch_model.bin", map_location="cpu")
model.load_state_dict(state_dict)
model.eval()

# Generate text
prompt = "Once upon a time,"
input_ids = torch.tensor([tokenizer.encode(prompt).ids])

with torch.no_grad():
    output = model.generate(
        input_ids,
        max_new_tokens=150,
        temperature=0.5,
        top_k=50,
        top_p=0.95
    )

print(tokenizer.decode(output[0].tolist()))

---

📄 License

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