MiniMind：3塊錢成本 + 2小時！訓練自己的0.02B的大模型。minimind源碼解讀、MOE架構

大家好，我是此林。

1. 前言

2. minimind模型源碼解讀

1.?MiniMind Config部分

1.1. 基礎參數

1.2. MOE配置

2.??MiniMind Model 部分?

2.1. MiniMindForCausalLM:?用于語言建模任務

2.2.?主干模型?MiniMindModel

2.3.?MiniMindBlock:?模型的基本構建塊

2.4.?class Attention(nn.Module)

2.5. MOEFeedForward

2.6. FeedForward

2.7. 其他

3. 寫在最后

1. 前言

大模型在這個時代可以說無處不在了，但是LLM動輒數百億參數的龐大規模。對于我們個人開發者而言不僅難以訓練，甚至連部署都顯得遙不可及。

那 github 上 20k Star+ 的開源項目 minimind，聲稱僅用3塊錢成本 + 2小時！即可訓練出僅為25.8M的超小語言模型?MiniMind。?

這不是謠言，此林已經幫你們試過了，AutoDL租用的 GPU 上訓練（Pretrain + SFT有監督微調）差不多2個小時半。Pretain 和 SFT 的數據集總共才 2.5G 左右。部分源碼也解讀了一下。

這是 Pretain 之后的：模型復讀機現象有點明顯，只會詞語接龍。

這是 SFT 微調后的：幻覺現象還是有點嚴重的，不過句子很流暢，可以接話了。

當然也用 react 快速搭了一個聊天框架，適配 http stream，看著還不錯。

鏈接文末自取。

2. minimind模型源碼解讀

1.?MiniMind Config部分

1.1. 基礎參數

from transformers import PretrainedConfig

PretrainedConfig 是所有 Transformer 模型配置類的基類，用于：

?? 存儲模型的結構參數（如隱藏層大小、注意力頭數、層數等）

?? 記錄預訓練模型的元信息（如 model_type、tokenizer_class）

?? 支持從預訓練模型自動加載配置

在 Transformers 中，每個模型都有一個對應的 Config 類，比如：

?? ???? ?BertModel — BertConfig

?? ???? ?GPT2Model — GPT2Config

?? ???? ?LlamaModel — LlamaConfig

這些都繼承自 PretrainedConfig，主要是構建模型前先配置參數。

使用場景舉例：查看 bert 隱藏層維度

from transformers import PretrainedConfigconfig = PretrainedConfig.from_pretrained("bert-base-uncased")print(config.hidden_size) # 查看隱藏層維度

那在這里的場景，是自定義配置用于訓練或推理，下面會說到。

這里就是定義了?MiniMindConfig，繼承自?PretrainedConfig。

class MiniMindConfig(PretrainedConfig):model_type = "minimind"def __init__(self,dropout: float = 0.0,bos_token_id: int = 1,# 省略...**kwargs):super().__init__(**kwargs)self.dropout = dropoutself.bos_token_id = bos_token_idself.eos_token_id = eos_token_idself.hidden_act = hidden_actself.hidden_size = hidden_size# 省略...# 和MoE相關的參數，如果use_moe=false，就忽略下邊的self.use_moe = use_moeself.num_experts_per_tok = num_experts_per_tok  # 每個token選擇的專家數量self.n_routed_experts = n_routed_experts  # 總的專家數量self.n_shared_experts = n_shared_experts  # 共享專家self.scoring_func = scoring_func  # 評分函數，默認為'softmax'self.aux_loss_alpha = aux_loss_alpha  # 輔助損失的alpha參數self.seq_aux = seq_aux  # 是否在序列級別上計算輔助損失self.norm_topk_prob = norm_topk_prob  # 是否標準化top-k概率

下面就一行一行來看里面的參數。

dropout: float = 0.0

Dropout 是一種防止過擬合的正則化技術，在每次前向傳播時，隨機丟棄一部分神經元的輸出，防止模型過度依賴某些神經元，從而提高泛化能力。

比如：dropout = 0.1，那么：

?? 模型每層中有 10% 的神經元在訓練時會被“屏蔽”（不參與計算）。

?? 在推理時（即模型使用階段），Dropout 是自動關閉的。

dropout: float = 0.0 就是關閉dropout。

bos_token_id: int = 1, # 開始 token 的 ID（Begin of Sentence）
eos_token_id: int = 2, # 結束 token 的 ID（End of Sentence）

在推理中的作用：

?? bos_token_id：模型知道“從這里開始生成”。

?? eos_token_id：模型在生成過程中，一旦預測出這個 token，就認為輸出完畢，停止生成。

這兩個 token 也經常用于?Hugging Face 的 generate() 方法

model.generate(
????input_ids,
????bos_token_id=1,
????eos_token_id=2,
????max_length=50
)

hidden_act: str = 'silu'

激活函數類型（如 silu、relu、gelu)，用SwiGLU激活函數替代了ReLU，這樣做是為了提高性能。

 hidden_size: int = 512

Transformer 每層的隱藏狀態維度?

intermediate_size: int = None

前饋層中間維度，如果為None，即用戶沒設置，后面代碼里會設置成 hidden_size * 8 / 3，這是在FeedForward里做的事情。

num_attention_heads: int = 8, # 每層中注意力頭的數量
num_hidden_layers: int = 8, # Transformer 層的數量
num_key_value_heads: int = 2, # KV heads 的數量（用于多頭注意力鍵值共享/分離）

每個 Transformer 層由多頭注意力層（Multi-Head Attention）和?前饋網絡（FFN）組成。

上面的參數表示：模型有 8 層 Transformer 層，每個 Transformer 層中有 8 個注意力頭，并且使用 2 個專門的頭來處理鍵（Key）和值（Value），相當于在多頭注意力的計算中，鍵和值部分的處理是分開的。

vocab_size: int = 6400

模型詞表大小（tokenizer 中的 token 總數）?，minimind是自己訓練了個tokenizer。? ? ? ??

rms_norm_eps: float = 1e-05, # RMSNorm 的 epsilon 值（防止除以0）rope_theta: int = 1000000.0, # RoPE 中的位置旋轉頻率（控制編碼的尺度）

采用了GPT-3的預標準化方法，也就是在每個Transformer子層的輸入上進行歸一化，而不是在輸出上。具體來說，使用的是RMSNorm歸一化函數。

像GPT-Neo一樣，去掉了絕對位置嵌入，改用了旋轉位置嵌入（RoPE），這樣在處理超出訓練長度的推理時效果更好。?

flash_attn: bool = True

Transformer 架構中，注意力機制?是關鍵的計算部分。傳統的注意力計算涉及較大的矩陣乘法，內存消耗大且計算速度較慢，尤其是在處理長序列時。FlashAttention 是一種基于 GPU 的優化算法，專門為 Transformer 模型中的自注意力計算（Self-Attention）進行加速。

1.2. MOE配置

下面的為MOE 配置：僅當 use_moe=True 時有效

它的結構基于Llama3和Deepseek-V2/3中的MixFFN混合專家模塊。

DeepSeek-V2在前饋網絡（FFN）方面，采用了更細粒度的專家分割和共享的專家隔離技術，以提高Experts的效果。

? ? ? ? ? ? use_moe: bool = False, # 是否啟用 MOE（專家混合）機制

? ? ? ? ? ? num_experts_per_tok: int = 2, # 每個 token 選擇的專家數量（Top-K）

? ? ? ? ? ? n_routed_experts: int = 4, # 可路由的專家總數（不包含共享專家）

? ? ? ? ? ? n_shared_experts: int = 1, # 所有 token 共享的專家數量（共享路徑）

? ? ? ? ? ? scoring_func: str = 'softmax', # 路由函數類型（如 softmax、top-k-gumbel）

? ? ? ? ? ? aux_loss_alpha: float = 0.1, # MOE 的輔助損失系數（平衡 load balance）

? ? ? ? ? ? seq_aux: bool = True, # 是否在序列級別計算輔助損失，而非 token 級別

? ? ? ? ? ? norm_topk_prob: bool = True, # 是否對 Top-K 路由概率歸一化（影響路由輸出

num_experts_per_tok: int = 2

在 MoE 中，我們通常有很多個前饋網絡（專家），比如 n_routed_experts = 8。但我們并不希望每個 token 都經過所有 8 個專家計算，這樣計算成本太高。所以我們使用一個門控網絡（gate）來為每個 token 選擇得分Top-K的專家處理。

這里等于 num_experts_per_tok = 2 就是選擇得分前 2 的專家，輸出結果是這兩個專家的加權和（按照 gate 輸出的概率加權）。

n_routed_experts: int = 4

n_routed_experts 表示有多少個普通專家（非共享專家）可以被 gate 路由（選擇）使用。

共享專家是指在所有層中都可以使用的專家，在token前向路徑自動經過，不用 gate 選。

結合剛才的?num_experts_per_tok = 2

對于每個 token：

gate 只會在這 4 個專家中計算得分（不是在全部 MoE 中的專家）。
從中選擇得分最高的 2 個來執行前饋計算。
gate 輸出加權這些專家的結果。

 scoring_func: str = 'softmax' # 路由函數類型（如 softmax、top-k-gumbel）

這個就是門控網絡（gate）對專家打分的機制，用了softmax，通常配合負載平衡機制使用，下面這個參數就是。

aux_loss_alpha: float = 0.1

在 MoE 模型中，每個 token 會通過路由選擇若干個專家來處理，這些專家的計算量通常是不均衡的，某些專家可能會頻繁被選中，而其他專家可能很少或幾乎不被選擇。這種不均衡的負載分配會導致一些專家被過度使用，而其他專家則被閑置，進而影響訓練效率和最終模型的泛化能力。

為了解決這個問題，輔助損失會通過在模型訓練中加上一個額外的損失項，強制使各個專家的使用頻率更均衡，從而改善負載平衡。

seq_aux: bool = True, # 是否在序列級別計算輔助損失，而非 token 級別

表示在序列級別計算輔助損失，而不是每個 token 單獨的負載。也就是模型會保證整個輸入序列的專家負載是均衡的，而不單單是某個具體的 token。在 token 級別計算輔助損失會導致較高的計算成本。

norm_topk_prob: bool = True

是否對 Top-K 路由概率歸一化（影響路由輸出），歸一化簡單來說就是讓概率總和為 1。

看到這里，相信你對MoE已經有了一定了解。

所以總的來說，MoE?模型的核心思想是：在每次前向傳播的過程中，通過門控網絡（gate）?只挑選得分Top-N個專家 參與計算，避免了全局計算的高成本。MoE 的最大優勢在于它的 稀疏性。

傳統的神經網絡是?Dense（密集）網絡，也叫 全連接網絡。對于每一個輸入樣本，網絡中的每個神經元都會參與計算，也就是每一層都會進行全量計算。每然后進行加權和計算。

?

2.??MiniMind Model 部分?

主要架構分為幾個部分，逐個來介紹。

2.1. MiniMindForCausalLM:?用于語言建模任務

包含：

1. 主干模型?MiniMindModel
2. 輸出層?lm_head（共享詞嵌入權重）
輸出：包含?logits（預測）、past_key_values（KV緩存）和?aux_loss（MOE專用）

2.2.?主干模型?MiniMindModel

包含：

1. 詞嵌入（embed_tokens）+?Dropout
2. 位置編碼（RoPE）預計算并注冊為?buffer
3. 堆疊多個?MiniMindBlock?層（用?nn.ModuleList）
輸出：最后的隱藏狀態、present?key?values、輔助損失（如果用了?MOE）

2.3.?MiniMindBlock:?模型的基本構建塊

包含：
1. 自注意力層（self_attn）
2. 兩個?RMSNorm?層（輸入?&?Attention?之后）
3. 一個前饋層（MLP?或?MOE）
4. 前向傳播：LayerNorm?→?Attention?→?殘差?→?LayerNorm?→?MLP?或?MOE→?殘差

self.mlp = FeedForward(config) if not config.use_moe else MOEFeedForward(config)

具體看這行代碼，如果use_moe == true，那么使用MOEFeedForward，否則使用FeedForward。