🧠 MiniMind

From-Scratch LLM Training Framework

Sub-modules

Python

Language

Files

~3.9K

Lines

🔵 Transformer · GQA 🟢 MoE · LoRA 🟡 RLHF · DPO · GRPO 🔴 26M–145M Params

💡 What is this?

MiniMind is a minimalist open-source project that lets you train a GPT-style language model entirely from scratch. Think of it as a teaching hospital for LLMs: small enough to understand every layer, yet complete enough to cover the full training lifecycle — from BPE tokenizer training through pretraining, supervised fine-tuning (SFT), LoRA adaptation, and reinforcement learning from human feedback (DPO, PPO, GRPO, SPO). The smallest config (26M params) trains on a single consumer GPU, making it ideal for learning how transformers really work.

Architecture Overview

The system is organised into four sub-modules. model/ defines the transformer backbone, trainer/ orchestrates every training recipe, dataset/ provides data loaders, and scripts/ handles deployment and inference.

flowchart LR
  subgraph model["model/"]
    MC[MiniMindConfig]
    MM[MiniMindModel
RoPE + GQA]
    MCL[MiniMindForCausalLM
GenerationMixin]
    ATT[Attention
Flash + KV Cache]
    FF[FeedForward
SwiGLU]
    MOE[MOEFeedForward
MoEGate + SharedExpert]
    RMS[RMSNorm]
    LORA[LoRA Module]
    MC --> MM
    MM --> MCL
    ATT --> MM
    FF --> MM
    MOE --> MM
    RMS --> MM
    LORA -.-> MCL
  end
  subgraph trainer["trainer/"]
    PT[pretrain.py]
    SFT[train_sft.py]
    TLORA[train_lora.py]
    DPO[train_dpo.py]
    PPO[train_ppo.py]
    GRPO[train_grpo.py]
    SPO[train_spo.py]
    KD[train_distillation.py]
    TR[train_reasoning.py]
    TOK[train_tokenizer.py]
    TU[trainer_utils.py]
  end
  subgraph dataset["dataset/"]
    PD[PretrainDataset]
    SD[SFTDataset]
    DD[DPODataset]
    RD[RLAIFDataset]
  end
  subgraph scripts["scripts/"]
    API[serve_api.py
FastAPI]
    WEB[web_demo.py
Streamlit]
    CVT[convert_model.py]
    CLI[api_client.py]
  end
  dataset --> trainer
  model --> trainer
  model --> scripts

Training Data Flow

Every training recipe follows a similar pipeline. Raw text is tokenised, batched by the dataset classes, passed through the model, and gradients flow back through the optimizer.

sequenceDiagram
    participant T as Tokenizer
(BPE, vocab=6400)
    participant D as Dataset
(PretrainDataset / SFTDataset)
    participant M as MiniMindForCausalLM
(Transformer + MoE)
    participant L as Loss
(CE / DPO / PPO)
    participant O as Optimizer
(AdamW)
    T->>D: Encode raw text → token IDs
    D->>M: Batched input_ids + labels
    M->>M: Embedding → N×(RMSNorm → Attention → RMSNorm → FFN) → LM Head
    M->>L: logits + labels → loss
    L->>O: loss.backward()
    O->>M: parameter update (lr=5e-4, weight_decay=0.1)
    Note over M,L: MoE adds aux_loss for load balancing

Narrated Walkthrough

🚶 Following the Full Training Pipeline

Step 1 — Tokenizer Training (trainer/train_tokenizer.py): Train a BPE tokenizer on your corpus to produce a 6,400-token vocabulary. This creates the tokenizer.model file consumed by all downstream stages.

Step 2 — Pretraining (trainer/pretrain.py): Feed raw text through PretrainDataset into MiniMindForCausalLM with standard cross-entropy loss. The model learns general language patterns over ~10B tokens. Distributed training via DDP is supported.

Step 3 — Supervised Fine-Tuning (trainer/train_sft.py): Switch to SFTDataset with instruction/response pairs. The loss mask ignores prompt tokens so the model only learns to generate responses.

Step 4 — LoRA Adaptation (trainer/train_lora.py): Freeze the base model, inject LoRA low-rank matrices into Q/V projections, and fine-tune on domain data with ~0.1% trainable params.

Step 5 — Preference Alignment (trainer/train_dpo.py): Use DPODataset to learn from chosen/rejected response pairs. DPO directly optimises the policy without a separate reward model.

Step 6 — GRPO Reinforcement (trainer/train_grpo.py): Group Relative Policy Optimization samples multiple completions per prompt, ranks them, and updates the policy using relative advantages — no critic network needed.

Step 7 — Deployment (scripts/serve_api.py): Wrap the trained model in an OpenAI-compatible FastAPI server. The scripts/web_demo.py Streamlit app provides a chat interface for interactive testing.

Key Components

Class / Function	File	Role
`MiniMindConfig`	model/model_minimind.py	HuggingFace-compatible config holding hidden_size, num_layers, MoE params, RoPE theta
`MiniMindModel`	model/model_minimind.py	Core transformer stack: token embedding → N × MiniMindBlock → RMSNorm
`MiniMindForCausalLM`	model/model_minimind.py	Full causal LM wrapping MiniMindModel + LM head + loss; inherits GenerationMixin
`Attention`	model/model_minimind.py	Grouped Query Attention (GQA) with KV cache, RoPE, optional Flash Attention
`FeedForward`	model/model_minimind.py	SwiGLU feed-forward: gate_proj · up_proj → SiLU → down_proj
`MoEGate`	model/model_minimind.py	Top-k softmax gating with auxiliary load-balancing loss
`MOEFeedForward`	model/model_minimind.py	Mixture of Experts layer dispatching tokens to top-k experts + shared expert
`RMSNorm`	model/model_minimind.py	Root Mean Square Layer Normalization (no centering, faster than LayerNorm)
`LoRA`	model/model_lora.py	Low-Rank Adaptation module injected into attention Q/V projections
`PretrainDataset`	dataset/lm_dataset.py	Streams tokenised raw-text chunks for next-token pretraining
`SFTDataset`	dataset/lm_dataset.py	Instruction-response pairs with prompt masking for supervised fine-tuning
`DPODataset`	dataset/lm_dataset.py	Chosen/rejected pairs for Direct Preference Optimization
`RLAIFDataset`	dataset/lm_dataset.py	Prompt dataset for RL-based training (PPO, GRPO, SPO)
`CriticModel`	trainer/train_ppo.py	PPO value network extending MiniMindForCausalLM with a scalar value head
`AutoAdaptiveValueTracker`	trainer/train_spo.py	SPO adaptive baseline tracker for reward normalisation

Deep-Dive Topics

Focused explorations of key architectural decisions and training strategies.

🏗️ Transformer Architecture

RoPE positional encoding, Grouped Query Attention with KV cache, RMSNorm, and SwiGLU feed-forward — how MiniMind assembles a GPT-style decoder.

model/model_minimind.py

🔄 Training Pipeline Flow

From pretrain through SFT, DPO, PPO, GRPO, and SPO — the complete lifecycle of training a language model with MiniMind.

trainer/*.py

🔀 MoE Routing

MoEGate softmax scoring, top-k expert selection, shared expert path, and auxiliary load-balancing loss for stable training.

model/model_minimind.py → MoEGate

🎯 LoRA Fine-tuning

Low-rank decomposition of weight updates, injection into Q/V projections, rank selection, and merging for efficient adaptation.

model/model_lora.py → LoRA

Directory Map

minimind/ ├── model/ │ ├── __init__.py │ ├── model_minimind.py # MiniMindConfig, MiniMindModel, MiniMindForCausalLM, Attention, FeedForward, MoEGate, MOEFeedForward, RMSNorm │ └── model_lora.py # LoRA module, apply_lora(), merge_lora() ├── trainer/ │ ├── pretrain.py # Pretraining loop (DDP, gradient accumulation) │ ├── train_sft.py # Supervised fine-tuning │ ├── train_lora.py # LoRA fine-tuning pipeline │ ├── train_dpo.py # Direct Preference Optimization │ ├── train_ppo.py # PPO with CriticModel value network │ ├── train_grpo.py # Group Relative Policy Optimization │ ├── train_spo.py # Self-Play Optimization with AutoAdaptiveValueTracker │ ├── train_distillation.py # Knowledge distillation from teacher model │ ├── train_reasoning.py # Reasoning-focused fine-tuning │ ├── train_tokenizer.py # BPE tokenizer training (vocab_size=6400) │ └── trainer_utils.py # Shared utilities: lr schedulers, logging, checkpointing ├── dataset/ │ ├── __init__.py │ └── lm_dataset.py # PretrainDataset, SFTDataset, DPODataset, RLAIFDataset ├── scripts/ │ ├── serve_api.py # OpenAI-compatible FastAPI server │ ├── web_demo.py # Streamlit chat interface │ ├── convert_model.py # Model format converter (HF ↔ custom) │ └── api_client.py # API client for testing serve_api.py └── eval_llm.py # Evaluation harness (imports model + trainer_utils)

Model Configuration (MiniMindConfig)

Key knobs exposed via MiniMindConfig(PretrainedConfig). The three standard model sizes are Small-26M, MoE-145M, and Base-104M.

Parameter	Small-26M	MoE-145M	Base-104M	Description
`hidden_size`	512	640	768	Transformer hidden dimension
`num_hidden_layers`	8	8	16	Number of transformer blocks
`num_attention_heads`	8	8	8	Query heads in multi-head attention
`num_key_value_heads`	2	2	2	KV heads for Grouped Query Attention
`vocab_size`	6400			BPE tokenizer vocabulary size
`max_position_embeddings`	32768			Maximum sequence length (with YaRN RoPE scaling)
`rope_theta`	1000000.0			RoPE base frequency for long-context extrapolation
`use_moe`	false	true	false	Enable Mixture of Experts layers
`num_experts`	—	8	—	Total expert count per MoE layer
`num_experts_per_tok`	—	2	—	Top-k experts activated per token
`n_shared_experts`	—	1	—	Always-on shared experts alongside routed ones
`flash_attn`	true (if available)			Use Flash Attention 2 for memory-efficient attention

Sub-Module Index

Drill down into each sub-module for class-level documentation, diagrams, and code walkthroughs.

Cross-Module Dependencies

flowchart TD
  dataset["dataset/
lm_dataset.py"]
  model_mm["model/
model_minimind.py"]
  model_lora["model/
model_lora.py"]
  trainer["trainer/*
(11 training scripts)"]
  scripts["scripts/*
(API, demo, converter)"]
  eval["eval_llm.py"]
  utils["trainer/
trainer_utils.py"]

  dataset -->|"PretrainDataset, SFTDataset, DPODataset"| trainer
  model_mm -->|"MiniMindForCausalLM"| trainer
  model_mm -->|"MiniMindForCausalLM"| scripts
  model_lora -->|"LoRA, apply_lora"| scripts
  model_mm -->|"MiniMindForCausalLM"| eval
  model_lora -->|"LoRA"| eval
  utils -->|"lr_scheduler, logging"| eval

  style dataset fill:#1a3a2a,stroke:#3fb950,color:#c9d1d9
  style model_mm fill:#1a2a3a,stroke:#58a6ff,color:#c9d1d9
  style model_lora fill:#1a2a3a,stroke:#58a6ff,color:#c9d1d9
  style trainer fill:#3a2a1a,stroke:#d29922,color:#c9d1d9
  style scripts fill:#3a1a1a,stroke:#f85149,color:#c9d1d9