Model Architecture

Transformer, MoE, and LoRA Implementation

📦 model/ 🐍 Python 🔖 83e52f6 📅 2026-03-26

⏱ Estimated reading time: 8 min

Files

Classes / Functions

~520

Lines of Code

Python

Language

🔵 Transformer · GQA 🟢 RoPE · YaRN 🟡 MoE · Top-K Gating 🔴 LoRA · Low-Rank

🏭 Think of it like a factory assembly line. The model/ directory is the blueprint for MiniMind's neural "factory." Raw token IDs enter the conveyor belt (embed_tokens), pass through a series of identical processing stations (MiniMindBlock), each of which inspects context (Attention) and refines features (FeedForward). At the end of the line, a quality-control head (lm_head) stamps out the next-token prediction. The MoE option swaps a single FFN station for a team of specialist workers, each activated only when their expertise is needed. LoRA lets you fine-tune the factory by clipping small adapter jigs onto existing machinery — no need to retool the whole line.

1 · Architecture Overview

MiniMind follows a decoder-only Transformer design compatible with Hugging Face's PreTrainedModel. The class hierarchy reflects clean separation of concerns: configuration lives in MiniMindConfig, the backbone in MiniMindModel, and the language-modelling head in MiniMindForCausalLM.

🔍 Click to zoom

classDiagram
    class MiniMindForCausalLM {
        +model: MiniMindModel
        +lm_head: Linear
        +forward(input_ids, labels)
        +weight tying
    }
    class MiniMindModel {
        +embed_tokens: Embedding
        +layers: List~MiniMindBlock~
        +norm: RMSNorm
        +freqs_cos, freqs_sin
    }
    class MiniMindBlock {
        +attention: Attention
        +feed_forward: FFN or MoE
        +attention_norm: RMSNorm
        +ffn_norm: RMSNorm
    }
    class Attention {
        +q_proj, k_proj, v_proj, o_proj
        +GQA: num_kv_heads=2
        +flash_attn support
        +KV cache
    }
    class FeedForward {
        +gate_proj: Linear
        +up_proj: Linear
        +down_proj: Linear
        +SwiGLU activation
    }
    class MOEFeedForward {
        +gate: MoEGate
        +experts: List~FeedForward~
        +shared_experts: FeedForward
    }
    class MoEGate {
        +weight: Linear
        +topk scoring
        +aux_loss
    }
    class RMSNorm {
        +weight: Parameter
        +eps: float
    }
    class MiniMindConfig {
        +hidden_size, num_layers
        +use_moe, n_routed_experts
        +rope_theta, flash_attn
    }
    class LoRA {
        +A: Linear (rank)
        +B: Linear (rank)
    }
    MiniMindForCausalLM --> MiniMindModel
    MiniMindModel --> MiniMindBlock
    MiniMindBlock --> Attention
    MiniMindBlock --> FeedForward
    MiniMindBlock --> MOEFeedForward
    MOEFeedForward --> MoEGate
    MOEFeedForward --> FeedForward : experts
    MiniMindBlock --> RMSNorm
    MiniMindForCausalLM ..> MiniMindConfig
    LoRA ..> Attention : adapts projections

2 · Key Classes & Functions

#	Name	File	Description
1	`MiniMindConfig`	model_minimind.py	Dataclass holding all hyperparameters; inherits `PretrainedConfig`.
2	`RMSNorm`	model_minimind.py	Root-mean-square layer normalisation with learnable scale.
3	`precompute_freqs_cis`	model_minimind.py	Builds cosine/sine tables for RoPE; supports YaRN scaling.
4	`apply_rotary_pos_emb`	model_minimind.py	Rotates Q/K tensors using precomputed frequencies.
5	`repeat_kv`	model_minimind.py	Repeats K/V heads to match Q head count in GQA.
6	`Attention`	model_minimind.py	Multi-head attention with grouped-query support and KV cache.
7	`FeedForward`	model_minimind.py	SwiGLU feed-forward network with auto-computed hidden dim.
8	`MoEGate`	model_minimind.py	Top-K softmax gating with auxiliary load-balancing loss.
9	`MOEFeedForward`	model_minimind.py	Mixture-of-Experts FFN with routed + shared experts.
10	`MiniMindBlock`	model_minimind.py	One Transformer layer: pre-norm → attention → FFN with residuals.
11	`MiniMindModel`	model_minimind.py	Full decoder backbone: embedding + N blocks + final norm.
12	`MiniMindForCausalLM`	model_minimind.py	Causal LM wrapper: backbone + lm_head + loss + generation.
13	`LoRA`	model_lora.py	Low-rank adapter module (A·B decomposition).
14	`apply_lora`	model_lora.py	Injects LoRA adapters into all square Linear layers.
15	`load_lora / save_lora`	model_lora.py	Serialization helpers for LoRA-only checkpoints.

3 · Forward-Pass Data Flow

The following diagram traces a single forward pass through MiniMindForCausalLM, from raw token IDs to cross-entropy loss.

🔍 Click to zoom

sequenceDiagram
    participant Caller as Trainer / generate()
    participant CLM as MiniMindForCausalLM
    participant Model as MiniMindModel
    participant Embed as embed_tokens
    participant Block as MiniMindBlock ×N
    participant Attn as Attention (GQA)
    participant FFN as FeedForward / MoE
    participant Norm as RMSNorm (final)
    participant Head as lm_head

    Caller->>CLM: input_ids, labels, past_kv
    CLM->>Model: input_ids, attention_mask, past_kv
    Model->>Embed: input_ids
    Embed-->>Model: hidden_states (B,T,D)
    loop For each layer i = 0..N-1
        Model->>Block: hidden_states, freqs, mask, kv_cache[i]
        Block->>Attn: RMSNorm(hidden_states)
        Attn-->>Block: attn_out, updated kv_cache
        Block->>Block: hidden_states += attn_out
        Block->>FFN: RMSNorm(hidden_states)
        FFN-->>Block: ffn_out [+ aux_loss if MoE]
        Block-->>Model: hidden_states += ffn_out
    end
    Model->>Norm: hidden_states
    Norm-->>Model: normed_states
    Model-->>CLM: normed_states, kv_caches, aux_loss
    CLM->>Head: normed_states[:, -logits_to_keep:]
    Head-->>CLM: logits (B,T,V)
    CLM->>CLM: CrossEntropyLoss + aux_loss
    CLM-->>Caller: CausalLMOutputWithPast

4 · Module Interactions

How model/ connects with the rest of the MiniMind project:

External Module	Direction	Description
`train_pretrain.py`	→ model	Instantiates `MiniMindForCausalLM`, calls `forward()`, reads loss + aux_loss.
`train_full_sft.py`	→ model	Fine-tunes full model weights; same forward path as pre-training.
`train_lora.py`	→ model, lora	Calls `apply_lora(model)`, trains only LoRA params, uses `save_lora()`.
`train_grpo.py`	→ model	Uses `logits_to_keep` param for memory-efficient GRPO training.
`dataset.py`	← model	Provides tokenized `input_ids` and `labels` consumed by the model.
`eval_model.py`	→ model	Loads checkpoint, calls `model.generate()` via `GenerationMixin`.
Hugging Face Hub	↔ model	Config + weights saved/loaded via `PreTrainedModel` API.

5 · Deep-Dive Hub

🔵

Attention Mechanism

Grouped-Query Attention with 8 Q-heads and 2 KV-heads. Supports Flash Attention via scaled_dot_product_attention and autoregressive KV caching for fast generation.

🟡

MoE Routing

Top-2 gating over 4 routed experts plus 1 shared expert. Softmax scoring with Kaiming-uniform init. Auxiliary load-balancing loss at token or sequence level keeps experts evenly utilised.

🟢

RoPE & YaRN

Rotary position embeddings with θ=1e6 and max 32768 positions. YaRN scaling (factor 16, β_fast=32, β_slow=1) for long-context extrapolation: f'(i) = f(i)·((1−γ) + γ/s).

🔴

LoRA Adapter

Rank-r decomposition injected into all square projection layers. A-matrix Gaussian-init, B-matrix zero-init ensures no change at start. Only ~0.5% extra parameters.

6 · Directory Map

model/ ├── __init__.py # Package init (empty) ├── model_minimind.py # Core transformer (~464 lines) │ ├── MiniMindConfig # Configuration dataclass │ ├── RMSNorm # Layer normalization │ ├── precompute_freqs_cis # RoPE frequency tables │ ├── apply_rotary_pos_emb # Rotary embedding application │ ├── repeat_kv # GQA key/value repetition │ ├── Attention # Multi-head GQA attention │ ├── FeedForward # SwiGLU FFN │ ├── MoEGate # Top-K expert gating │ ├── MOEFeedForward # MoE layer │ ├── MiniMindBlock # Transformer block │ ├── MiniMindModel # Full decoder backbone │ └── MiniMindForCausalLM # Causal LM head + loss └── model_lora.py # LoRA adapters (~54 lines) ├── LoRA # Low-rank module ├── apply_lora # Injection helper ├── load_lora # Weight loading └── save_lora # Weight saving

7 · Configuration Reference

All knobs live in MiniMindConfig(PretrainedConfig). Values shown are the defaults.

Parameter	Default	Description
`hidden_size`	512 / 640 / 768	Embedding dimension (varies by model size).
`num_hidden_layers`	8 / 16	Number of Transformer blocks.
`num_attention_heads`	8	Number of query heads per layer.
`num_key_value_heads`	2	Number of KV heads (grouped-query attention).
`vocab_size`	6400	Vocabulary size used by the tokenizer.
`max_position_embeddings`	32768	Maximum sequence length supported.
`rope_theta`	1e6	Base frequency for rotary position embeddings.
`use_moe`	False	Enable Mixture-of-Experts feed-forward layers.
`num_experts_per_tok`	2	Top-K experts activated per token.
`n_routed_experts`	4	Total number of routed expert FFNs.
`n_shared_experts`	1	Number of always-active shared experts.
`scoring_func`	"softmax"	Gating score function for MoE routing.
`aux_loss_alpha`	0.01	Weight of auxiliary load-balancing loss.
`flash_attn`	True	Use PyTorch `scaled_dot_product_attention`.
`dropout`	0.0	Dropout probability (disabled by default).
`rms_norm_eps`	1e-5	Epsilon for RMSNorm numerical stability.
`hidden_act`	"silu"	Activation function (`ACT2FN` lookup).
`inference_rope_scaling`	YaRN dict	YaRN config: factor=16, beta_fast=32, beta_slow=1.

8 · Code Examples

8.1 · SwiGLU Feed-Forward (from `model_minimind.py`)

class FeedForward(nn.Module):
    def __init__(self, config: MiniMindConfig):
        super().__init__()
        hidden_dim = int(config.hidden_size * 8 / 3)
        hidden_dim = 64 * ((hidden_dim + 63) // 64)  # round to 64
        self.gate_proj = nn.Linear(config.hidden_size, hidden_dim, bias=False)
        self.up_proj   = nn.Linear(config.hidden_size, hidden_dim, bias=False)
        self.down_proj = nn.Linear(hidden_dim, config.hidden_size, bias=False)
        self.act_fn    = ACT2FN[config.hidden_act]

    def forward(self, x):
        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))

8.2 · LoRA Injection (from `model_lora.py`)

class LoRA(nn.Module):
    def __init__(self, in_features, out_features, rank=8):
        super().__init__()
        self.A = nn.Linear(in_features, rank, bias=False)
        self.B = nn.Linear(rank, out_features, bias=False)
        nn.init.normal_(self.A.weight, std=1/rank)
        nn.init.zeros_(self.B.weight)

    def forward(self, x):
        return self.B(self.A(x))

def apply_lora(model, rank=8):
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear) and module.in_features == module.out_features:
            lora = LoRA(module.in_features, module.out_features, rank)
            setattr(module, 'lora', lora)
            original_forward = module.forward
            module.forward = lambda x, of=original_forward, l=lora: of(x) + l(x)

8.3 · YaRN-scaled RoPE Frequency Computation (from `model_minimind.py`)

def precompute_freqs_cis(dim, end, theta=1e6, scaling_config=None):
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
    if scaling_config is not None:
        # YaRN: f'(i) = f(i) * ((1 - gamma) + gamma / scale)
        factor    = scaling_config["factor"]        # 16
        beta_fast = scaling_config["beta_fast"]      # 32
        beta_slow = scaling_config["beta_slow"]      # 1
        # Compute linear ramp gamma ∈ [0, 1]
        low  = math.floor(dim * math.log(end / (beta_fast * 2 * math.pi)) / (2 * math.log(theta)))
        high = math.ceil(dim * math.log(end / (beta_slow * 2 * math.pi)) / (2 * math.log(theta)))
        gamma = (torch.arange(0, dim // 2).float() - low) / (high - low)
        gamma = gamma.clamp(0, 1)
        freqs = freqs * ((1 - gamma) + gamma / factor)
    t = torch.arange(end, dtype=torch.float32)
    freqs = torch.outer(t, freqs)
    return torch.cos(freqs), torch.sin(freqs)