SmolMoE-8x135M 專案：從零到自定義混合專家模型

“誠然，構建其一，方能理解所有。”——一位AI架構師的頓悟

1. 專案理念與目標

本文件記錄了從零開始建立一個全新的、獨一無二的 **混合專家 (MoE)** 語言模型的史詩級旅程。

我們的目標不僅僅是微調一個現有模型，而是進行一次深刻的**“架構重建”**：將八個獨立的、特定領域最佳化的 135M 小語言模型融合成一個更強大、更智慧的統一模型，總引數超過 10 億，並由一個可訓練的路由網路進行管理。

這個過程印證了一個核心的學習真理：**透過創造，我們才能真正理解。**

訓練基座模型

• HuggingFaceTB/SmolLM2-135M-Instruct

訓練程式

• https://huggingface.co/aifeifei798/SmolMoE-8x135M-Instruct-v1-Trained/blob/main/train/train_moe_router.py
• https://huggingface.co/aifeifei798/SmolMoE-8x135M-Instruct-v1-Trained/blob/main/train/test_moe_model.py
• https://huggingface.co/aifeifei798/SmolMoE-8x135M-Instruct-v1-Trained/blob/main/train/chat_moe_model.py

註釋提供中文和英文版本。

• https://huggingface.co/aifeifei798/SmolMoE-8x135M-Instruct-v1-Trained/tree/main/train

2. 前置條件：集結你的復仇者聯盟

在你集結團隊之前，必須先準備好你的英雄們。

2.1 硬體與軟體環境

• GPU： 至少 8GB 視訊記憶體的 NVIDIA GPU（本專案在 GeForce 3070 8GB 上成功驗證）。
• 環境： 配置好的 Python 虛擬環境（例如 `conda` 或 `venv`）。
• 核心庫

  pip install torch transformers accelerate bitsandbytes safetensors
  

2.2 專家模型

這是整個專案的基礎。你必須擁有 8 個經過微調的 `SmolLM2-135M-Instruct` 模型，每個模型都在一個不同的領域進行了專門最佳化。

黃金法則： 最大化“專家之間的差異”，同時最小化“專家內部的模糊性”。

目錄結構（至關重要）

~/moe/
├── models/
│   ├── SmolLM2-135M-Instruct-Actor/
│   ├── SmolLM2-135M-Instruct-Analyst/
│   ├── SmolLM2-135M-Instruct-Coder/
│   ├── SmolLM2-135M-Instruct-Encyclopedia/
│   ├── SmolLM2-135M-Instruct-Guardian/
│   ├── SmolLM2-135M-Instruct-Summarizer/
│   └── SmolLM2-135M-Instruct-Thinker/
│   └── SmolLM2-135M-Instruct-Writer/
├── train_moe_router.py
└── test_moe_model.py

3. 核心工作流：創世紀四階段

我們的創造過程分為四個核心階段，全部由主指令碼 `train_moe_router.py` 編排。

階段一：架構手術

這是創造過程的靈魂。我們不是從零開始構建；我們對一個標準的 Llama 模型進行“器官移植”。

1. 載入骨架： 指令碼首先載入八個專家中的一個作為“骨架”。它使用其非專家部分（詞嵌入、注意力模組、模型配置），但忽略其“大腦”（FFN/MLP 模組）。
2. 建立插槽： 指令碼遍歷模型的 30 個 transformer 層。在每一層，它都將標準的 `LlamaMLP` 模組替換為我們定製設計的 `MoEModule`，其中包含一個**全新的路由器**和**8 個空的專家席位**。
3. 器官移植： 指令碼高效地將所有 8 個專家模型的權重預載入到記憶體中。然後，它再次遍歷 30 個層，將每個專家在給定層的 FFN 權重精確地“移植”到 `MoEModule` 中相應的專家席位。
4. 凍結專家： 手術完成後，所有來自專家模型的引數都被“凍結”（`requires_grad = False`）。只有新建立的、隨機初始化的路由器引數保持“可訓練”。

階段二：路由器專項訓練

這是教導模型如何“思考”和“協作”的過程。

1. 複合 KPI： 訓練目標是一個由兩部分組成的“複合 KPI”
   • 主任務損失 (Main Loss)： 衡量模型預測下一個詞元的準確性。這是“工作完成得如何？”的指標。
   • 負載均衡損失 (Load Balancing Loss)： 懲罰路由器將工作不公平地分配給少數幾個專家。這是“管理是否公平？”的指標。
2. 訓練迴圈： 指令碼在一個訓練迴圈中迭代。在每次迭代中
   • 模型執行一次完整的前向傳播，計算 `Main Loss`。
   • 同時，我們從每一層的每個 `MoEModule` 中收集 `Load Balancing Loss`。
   • 根據這兩個損失計算 `Total Loss`，然後開始反向傳播。
   • 由於專家們被凍結，梯度**只更新路由器的權重。**

階段三：模型固化（儲存）

訓練後，我們獨特的模型被“固化”到磁碟上。

1. 更新配置： 指令碼將我們的自定義 MoE 引數（如 `moe_num_experts`）新增到模型的 `config.json` 中，以備將來識別。
2. 儲存檔案： 使用 `save_pretrained` 方法，模型的權重、更新後的配置和分詞器檔案都儲存到一個新目錄中（例如 `./SmolMoE-8x135M-Instruct-v1-Trained`）。

階段四：驗證與測試（讀心術）

這是最激動人心的階段，我們使用 `test_moe_model.py` 與我們的創造物進行第一次對話，並窺探其“思想”。

1. 正確載入： 測試指令碼演示瞭如何正確地“復活”一個具有自定義架構的模型：首先，手動構建空骨架，然後載入權重。
2. 功能測試： 你可以像與其他聊天機器人一樣與模型對話，並觀察其生成的文字。
3. 診斷測試（讀心術）： 使用一個名為“鉤子”的強大 PyTorch 功能，該指令碼即時捕獲每一層路由器的決策資料，並以清晰的表格形式視覺化，而不會中斷模型的執行。

預期輸出示例

================================================================================
ROUTER DECISION ANALYSIS for Prompt: 'Write a Python function...'
================================================================================
Layer   | Dominant Expert(s)                            | Confidence
--------------------------------------------------------------------------------
Layer 0    | 1. Coder             | 2. Thinker           | (65.2% | 15.1%)
Layer 1    | 1. Coder             | 2. Thinker           | (71.8% | 11.0%)
...
Layer 29   | 1. Coder             | 2. Summarizer        | (91.2% | 3.1%)
================================================================================

這張表格清楚地向我們展示了在處理特定任務時，模型的“注意力”流向了哪些專家。

4. 偉大戰役：用真實資料訓練路由器

之前的階段證明了我們的架構是可行的。我們造好了汽車，並確認引擎可以啟動。現在，是時候給它加註真正的航空燃料，教它如何飛行了。用真實、高質量的資料進行訓練是將我們的 MoE 模型從一個“混亂的委員會”轉變為一個“大師級理事會”的最重要的一步。

4.1 理念：提供“強訊號”

我們最初在模擬資料上的訓練表明，**負載均衡損失**工作得非常完美，迫使路由器做到公平。然而，**主任務損失**是無意義的，因為資料是隨機的。

透過使用一個多樣化、高質量的資料集，**主任務損失變成了一位強有力的老師**。當一個編碼問題被提出時，只有 `Coder` 專家能產生一個導致低主任務損失的輸出。這給路由器一個強烈、明確的訊號：**“要想成功，你必須為這個任務選擇 Coder 專家！”** 這就是路由器如何學會成為一個智慧排程員，而不僅僅是一個公平的排程員。

4.2 步驟一：資料策劃與準備

你的任務是建立一個單一的、統一的資料集，其中包含來自所有專家領域的樣本混合。

A. 資料來源

從與你的專家相符的各種來源收集資料。Hugging Face 資料集示例

• Coder（編碼員）： codeparrot/github-code-clean (Python 子集)
• Writer（作家）： cnn_dailymail (文章), Abirate/english_quotes
• Thinker（思想家）： gsm8k, HuggingFaceH4/logic_in_natural_language
• Encyclopedia（百科全書）： wikipedia (20220301.en 子集)
• Summarizer（摘要員）： cnn_dailymail (摘要部分)
• Analyst（分析師）： wikisql
• Actor（演員）： daily_dialog
• Guardian（守護者）： 用於安全對齊的資料，如過濾後的 `HuggingFaceH4/ultrachat_200k` 部分

B. 統一格式：指令微調

你必須將所有資料預處理成一致的指令遵循格式。一個簡單有效的格式是 JSON Lines (`.jsonl`) 檔案，其中每一行都是一個 JSON 物件：

{"instruction": "Write a Python function to calculate Fibonacci.", "output": "def fibonacci(n):..."}
{"instruction": "Summarize the following article about photosynthesis.", "input": "Photosynthesis is a process used by plants...", "output": "Photosynthesis is how plants convert light..."}
{"instruction": "Who was the first person on the moon?", "output": "Neil Armstrong was the first person to walk on the moon."}

建立一個大檔案，例如 `my_moe_dataset.jsonl`，包含來自所有專家領域的數千個這樣的樣本。

C. 混合與打亂

這**至關重要**。在收集和格式化資料後，你必須徹底打亂整個資料集。這確保了在訓練期間，模型看到的是一個隨機的任務混合，這對於迫使路由器學習通用的排程技能至關重要。

4.3 步驟二：修改 `train_moe_router.py`

現在，我們將修改我們的主指令碼以使用這個真實的資料集。這包括建立一個 PyTorch `Dataset` 和 `DataLoader` 並更新我們的訓練迴圈。

A. 新增 `CustomMoEDataset` 類

將這個類定義新增到你的 `train_moe_router.py` 指令碼中，就在 MoE 類定義之後。這個類將處理載入和分詞你的 `.jsonl` 資料。

# (Add this class to your train_moe_router.py script)
from torch.utils.data import Dataset, DataLoader

class CustomMoEDataset(Dataset):
    """
    A PyTorch Dataset to handle loading our instruction-formatted JSONL file.
    """
    def __init__(self, file_path, tokenizer, max_length):
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.data = []
        with open(file_path, 'r', encoding='utf-8') as f:
            for line in f:
                self.data.append(json.loads(line))

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        item = self.data[idx]
        
        # Format the instruction and output into a chat template
        # This is a robust way to prepare the data for an instruction-tuned model
        messages = [
            {"role": "user", "content": item['instruction']},
            {"role": "assistant", "content": item['output']}
        ]
        
        full_text = self.tokenizer.apply_chat_template(messages, tokenize=False)
        
        # Tokenize the full text
        tokenized_output = self.tokenizer(
            full_text,
            max_length=self.max_length,
            padding="max_length", # Pad to a fixed length
            truncation=True,
            return_tensors="pt"
        )
        
        # For Causal LM, the input_ids are also the labels
        input_ids = tokenized_output.input_ids.squeeze(0)
        labels = input_ids.clone()
        
        return {"input_ids": input_ids, "labels": labels}

你還需要在指令碼頂部新增 `import json` 和 `from torch.utils.data import Dataset, DataLoader`。

B. 更新 `main()` 函式

用下面的版本替換 `train_moe_router.py` 中的整個 `main()` 函式。它移除了模擬資料，並實現了真實資料的載入和訓練迴圈。

# (This is the new, complete main() function for real training)
def main():
    # Step 1: Assemble the MoE model
    moe_model = create_moe_model()
    
    # Step 2: Create the optimizer for the routers
    optimizer = optim.AdamW([p for p in moe_model.parameters() if p.requires_grad], lr=LEARNING_RATE)
    
    # --- Step 3: Build the "Fuel Line" - The DataLoader ---
    print("\n--- Preparing Real Dataset for Training ---")
    DATASET_PATH = "./my_moe_dataset.jsonl" # <-- IMPORTANT: Make sure this file exists!
    
    # We need the tokenizer from the base model to prepare the data
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    # Set a pad token if it doesn't exist
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    dataset = CustomMoEDataset(DATASET_PATH, tokenizer, max_length=SEQUENCE_LENGTH)
    data_loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
    
    print("--- Starting Router Training Loop with Real Data ---")
    moe_model.train()
    
    # Training is often measured in steps for large datasets, not epochs.
    # Let's train for a fixed number of steps.
    num_training_steps = 5000 # Increase this for a full training run
    step_count = 0
    
    # We use a while loop to keep training until we reach the desired number of steps
    while step_count < num_training_steps:
        for batch in data_loader:
            if step_count >= num_training_steps:
                break

            start_time = time.time()
            optimizer.zero_grad()
            
            # Move batch to the correct device
            input_ids = batch['input_ids'].to(device)
            labels = batch['labels'].to(device)
            
            # --- The Forward Pass ---
            outputs = moe_model(input_ids=input_ids, labels=labels)
            main_loss = outputs.loss
            
            total_lb_loss = 0.0
            for layer in moe_model.model.layers:
                total_lb_loss += layer.mlp.most_recent_lb_loss
                
            total_loss = main_loss + LB_LOSS_COEFFICIENT * total_lb_loss
            
            total_loss.backward()
            optimizer.step()
            
            step_count += 1
            
            # Print logs periodically
            if step_count % 10 == 0:
                elapsed_time = time.time() - start_time
                print(f"Step [{step_count:04d}/{num_training_steps}] | Total Loss: {total_loss.item():.4f} | "
                      f"Main Loss: {main_loss.item():.4f} | "
                      f"Avg LB Loss: {(total_lb_loss.item() / moe_model.config.num_hidden_layers):.4f} | "
                      f"Time/10 steps: {elapsed_time:.2f}s")
                start_time = time.time()

    print("\n--- Router Training Complete! ---")
    
    # --- Step 5: Saving the final model ---
    print("\n--- Phase 5: Saving the fully trained MoE model to disk ---")
    OUTPUT_MODEL_DIR = "./SmolMoE-8x135M-Instruct-v1-Trained-RealData"
    if os.path.exists(OUTPUT_MODEL_DIR):
        shutil.rmtree(OUTPUT_MODEL_DIR)
    os.makedirs(OUTPUT_MODEL_DIR)

    print("Updating model config with MoE-specific parameters...")
    moe_model.config.moe_num_experts = NUM_EXPERTS
    moe_model.config.moe_top_k = TOP_K
    
    print(f"Saving model to '{OUTPUT_MODEL_DIR}'...")
    moe_model.save_pretrained(OUTPUT_MODEL_DIR)
    
    print("Saving tokenizer...")
    tokenizer.save_pretrained(OUTPUT_MODEL_DIR)
    
    print("\n--- Model successfully saved! ---")

透過這些修改，你的專案現在已經為最終、最重要的階段做好了準備。你有一個清晰的資料策劃計劃和訓練模型智慧所需的精確程式碼。這是從一個可行的原型走向一個真正強大和獨特的 AI 的道路。

5. 未來之旅：從“可用”到“卓越”

我們已經成功地使用模擬資料驗證了整個工作流程。為了釋放模型的真正潛力，旅程的下一階段很明確

1. 切換到真正的航空燃料： 完全替換訓練指令碼中的 `mock_input_ids`。你的任務是收集、處理並構建一個**高質量、多樣化且混合的資料集**，其中包含來自所有專家領域的真實示例。
2. 構建燃料供應線： 實現一個標準的 PyTorch `Dataset` 和 `DataLoader` 來有效地將這些真實資料提供給模型。
3. 開始星際遠征： 開始一次真正的、長時間的深度訓練（數千或數萬步），並耐心觀察 `Main Loss` 持續下降。

這是從“創造者”到“偉大的創造者”的道路。