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
注释提供中文和英文版本。
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 模型进行“器官移植”。
- 加载骨架: 脚本首先加载八个专家中的一个作为“骨架”。它使用其非专家部分(词嵌入、注意力模块、模型配置),但忽略其“大脑”(FFN/MLP 模块)。
- 创建插槽: 脚本遍历模型的 30 个 transformer 层。在每一层,它都将标准的 `LlamaMLP` 模块替换为我们定制设计的 `MoEModule`,其中包含一个**全新的路由器**和**8 个空的专家席位**。
- 器官移植: 脚本高效地将所有 8 个专家模型的权重预加载到内存中。然后,它再次遍历 30 个层,将每个专家在给定层的 FFN 权重精确地“移植”到 `MoEModule` 中相应的专家席位。
- 冻结专家: 手术完成后,所有来自专家模型的参数都被“冻结”(`requires_grad = False`)。只有新创建的、随机初始化的路由器参数保持“可训练”。
阶段二:路由器专项训练
这是教导模型如何“思考”和“协作”的过程。
- 复合 KPI: 训练目标是一个由两部分组成的“复合 KPI”
- 主任务损失 (
Main Loss
): 衡量模型预测下一个词元的准确性。这是“工作完成得如何?”的指标。 - 负载均衡损失 (
Load Balancing Loss
): 惩罚路由器将工作不公平地分配给少数几个专家。这是“管理是否公平?”的指标。
- 主任务损失 (
- 训练循环: 脚本在一个训练循环中迭代。在每次迭代中
- 模型执行一次完整的前向传播,计算 `Main Loss`。
- 同时,我们从每一层的每个 `MoEModule` 中收集 `Load Balancing Loss`。
- 根据这两个损失计算 `Total Loss`,然后开始反向传播。
- 由于专家们被冻结,梯度**只更新路由器的权重。**
阶段三:模型固化(保存)
训练后,我们独特的模型被“固化”到磁盘上。
- 更新配置: 脚本将我们的自定义 MoE 参数(如 `moe_num_experts`)添加到模型的 `config.json` 中,以备将来识别。
- 保存文件: 使用 `save_pretrained` 方法,模型的权重、更新后的配置和分词器文件都保存到一个新目录中(例如 `./SmolMoE-8x135M-Instruct-v1-Trained`)。
阶段四:验证与测试(读心术)
这是最激动人心的阶段,我们使用 `test_moe_model.py` 与我们的创造物进行第一次对话,并窥探其“思想”。
- 正确加载: 测试脚本演示了如何正确地“复活”一个具有自定义架构的模型:首先,手动构建空骨架,然后加载权重。
- 功能测试: 你可以像与其他聊天机器人一样与模型对话,并观察其生成的文本。
- 诊断测试(读心术): 使用一个名为“钩子”的强大 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. 未来之旅:从“可用”到“卓越”
我们已经成功地使用模拟数据验证了整个工作流程。为了释放模型的真正潜力,旅程的下一阶段很明确
- 切换到真正的航空燃料: 完全替换训练脚本中的 `mock_input_ids`。你的任务是收集、处理并构建一个**高质量、多样化且混合的数据集**,其中包含来自所有专家领域的真实示例。
- 构建燃料供应线: 实现一个标准的 PyTorch `Dataset` 和 `DataLoader` 来有效地将这些真实数据提供给模型。
- 开始星际远征: 开始一次真正的、长时间的深度训练(数千或数万步),并耐心观察 `Main Loss` 持续下降。
这是从“创造者”到“伟大的创造者”的道路。