Upload 7 files

train/chat_moe_model.py

@@ -0,0 +1,137 @@
+# ==============================================================================
+#                 Smol-MoE 8x135M - "Chat with Your Creation"
+#               (Final Interactive Inference Script)
+# ==============================================================================
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+from transformers.models.llama.modeling_llama import LlamaMLP
+import os
+
+# --- 1. 关键：重新定义你的所有自定义模块 ---
+# 这是让 from_pretrained() 能够成功重建你自定义模型的关键。
+
+MODEL_PATH = "./SmolMoE-8x135M-Instruct-v1-Trained"
+
+# 从保存好的模型配置中读取MoE参数
+config = AutoConfig.from_pretrained(MODEL_PATH)
+NUM_EXPERTS = config.moe_num_experts
+TOP_K = config.moe_top_k
+
+class MoERouter(nn.Module):
+    def __init__(self, hidden_size: int, num_experts: int):
+        super().__init__()
+        self.layer = nn.Linear(hidden_size, num_experts, bias=False)
+    def forward(self, hidden_states):
+        return self.layer(hidden_states)
+
+class MoEModule(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.top_k = TOP_K
+        self.num_experts = NUM_EXPERTS
+        self.router = MoERouter(self.hidden_size, self.num_experts)
+        self.experts = nn.ModuleList([LlamaMLP(config) for _ in range(self.num_experts)])
+
+    def forward(self, hidden_states):
+        original_shape = hidden_states.shape
+        flat_hidden_states = hidden_states.view(-1, self.hidden_size)
+        router_logits = self.router(flat_hidden_states)
+        routing_weights = F.softmax(router_logits, dim=-1, dtype=torch.float)
+        routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
+        routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
+        routing_weights = routing_weights.to(hidden_states.dtype)
+        final_hidden_states = torch.zeros_like(flat_hidden_states)
+        for k in range(self.top_k):
+            expert_indices_k = selected_experts[:, k]
+            routing_weights_k = routing_weights[:, k]
+            for i in range(self.num_experts):
+                mask = expert_indices_k == i
+                if mask.any():
+                    expert_output = self.experts[i](flat_hidden_states[mask])
+                    final_hidden_states.index_add_(0, torch.where(mask)[0], expert_output * routing_weights_k[mask].unsqueeze(1))
+        return final_hidden_states.view(*original_shape)
+
+
+# --- 2. 主程序：加载模型并开始对话 ---
+def main():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    # --- 模型加载 ---
+    print(f"Loading tokenizer from '{MODEL_PATH}'...")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+    
+    print(f"Manually rebuilding MoE model structure...")
+    # 用`from_config`创建一个随机权重的、但结构正确的“空壳”
+    moe_model = AutoModelForCausalLM.from_config(config)
+    
+    # 手动进行“架构手术”，把标准的MLP替换成我们的MoE模块
+    for i, layer in enumerate(moe_model.model.layers):
+        layer.mlp = MoEModule(config)
+        
+    print(f"Loading your trained MoE weights into the correct structure...")
+    from safetensors.torch import load_file
+    state_dict = load_file(os.path.join(MODEL_PATH, "model.safetensors"), device="cpu")
+    
+    # 使用`strict=False`灵活加载，然后手动绑定权重
+    moe_model.load_state_dict(state_dict, strict=False)
+    moe_model.tie_weights()
+    
+    moe_model.to(device, dtype=torch.bfloat16)
+    moe_model.eval() # 切换到评估模式
+    print("--- MoE Model is ready for conversation! ---")
+    print("Type 'exit' or 'quit' to end the chat.\n")
+
+    # --- 交互式对话循环 ---
+    messages = []
+    while True:
+        try:
+            user_input = input("You: ")
+            if user_input.lower() in ["exit", "quit"]:
+                print("Goodbye!")
+                break
+
+            # 1. 将用户输入添加到对话历史
+            messages.append({"role": "user", "content": user_input})
+            
+            # 2. 使用聊天模板格式化完整的对话历史
+            #   `add_generation_prompt=True` 会在末尾添加助手角色的起始标记
+            prompt_text = tokenizer.apply_chat_template(
+                messages, 
+                tokenize=False, 
+                add_generation_prompt=True
+            )
+            
+            # 3. 编码输入并发送到GPU
+            inputs = tokenizer(prompt_text, return_tensors="pt").to(device)
+
+            # 4. 生成回复
+            with torch.no_grad():
+                outputs = moe_model.generate(
+                    **inputs, 
+                    max_new_tokens=256, 
+                    temperature=0.7, 
+                    top_p=0.9, 
+                    do_sample=True
+                )
+            
+            # 5. 解码并清理输出
+            # `outputs[0]` 包含了完整的对话（输入+输出），我们需要提取出模型新生成的部分
+            full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+            # 通过移除原始prompt来找到新生成的部分
+            model_response = full_response.replace(prompt_text.replace("<s> ", "").replace("</s>", ""), "").strip()
+            
+            print(f"MoE Model: {model_response}")
+            
+            # 6. 将模型的回复也添加到对话历史中，以便进行多轮对话
+            messages.append({"role": "assistant", "content": model_response})
+
+        except KeyboardInterrupt:
+            print("\nGoodbye!")
+            break
+
+if __name__ == "__main__":
+    main()

train/chat_moe_model_zhcn_en.py

@@ -0,0 +1,177 @@
+# ==============================================================================
+#                 Smol-MoE 8x135M - "Chat with Your Creation"
+#               (Final Interactive Inference Script)
+#
+#                 Smol-MoE 8x135M - “与你的造物对话”
+#               (最终版交互式推理脚本)
+# ==============================================================================
+
+# --- Core Library Imports / 核心库导入 ---
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+from transformers.models.llama.modeling_llama import LlamaMLP
+import os
+
+# --- 1. CRITICAL: Re-define Your Custom Architecture ---
+# --- 1. 关键：重新定义你的所有自定义模块 ---
+# When loading a model with a custom architecture, Hugging Face needs to know the definition of the custom classes.
+# By defining them here, we allow the `from_pretrained` process to correctly reconstruct our unique MoE model.
+# 在加载一个拥有自定义架构的模型时，Hugging Face 需要知道这些自定义类的定义。
+# 在这里定义它们，我们才能让 `from_pretrained` 函数成功地重建我们独一-无二的MoE模型。
+
+# --- Model Configuration / 模型配置 ---
+MODEL_PATH = "./SmolMoE-8x135M-Instruct-v1-Trained"
+
+# Load our custom MoE parameters from the saved config file.
+# 从我们保存的配置文件中，加载自定义的MoE参数。
+config = AutoConfig.from_pretrained(MODEL_PATH)
+NUM_EXPERTS = config.moe_num_experts
+TOP_K = config.moe_top_k
+
+class MoERouter(nn.Module):
+    """The Router module. Its job is to score experts for each token."""
+    """路由器模块。它的工作是为每个token给所有专家打分。"""
+    def __init__(self, hidden_size: int, num_experts: int):
+        super().__init__()
+        self.layer = nn.Linear(hidden_size, num_experts, bias=False)
+    def forward(self, hidden_states):
+        return self.layer(hidden_states)
+
+class MoEModule(nn.Module):
+    """The custom Mixture-of-Experts module that replaces the standard FFN."""
+    """我们自定义的混合专家模块，它替换了标准的FFN。"""
+    def __init__(self, config):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.top_k = TOP_K
+        self.num_experts = NUM_EXPERTS
+        self.router = MoERouter(self.hidden_size, self.num_experts)
+        self.experts = nn.ModuleList([LlamaMLP(config) for _ in range(self.num_experts)])
+
+    def forward(self, hidden_states):
+        original_shape = hidden_states.shape
+        flat_hidden_states = hidden_states.view(-1, self.hidden_size)
+        router_logits = self.router(flat_hidden_states)
+        routing_weights = F.softmax(router_logits, dim=-1, dtype=torch.float)
+        routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
+        routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
+        routing_weights = routing_weights.to(hidden_states.dtype)
+        final_hidden_states = torch.zeros_like(flat_hidden_states)
+        for k in range(self.top_k):
+            expert_indices_k = selected_experts[:, k]
+            routing_weights_k = routing_weights[:, k]
+            for i in range(self.num_experts):
+                mask = expert_indices_k == i
+                if mask.any():
+                    expert_output = self.experts[i](flat_hidden_states[mask])
+                    final_hidden_states.index_add_(0, torch.where(mask)[0], expert_output * routing_weights_k[mask].unsqueeze(1))
+        return final_hidden_states.view(*original_shape)
+
+# --- 2. Main Program: Load Model and Start Conversation ---
+# --- 2. 主程序：加载模型并开始对话 ---
+def main():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    # --- Model Loading / 模型加载 ---
+    print(f"Loading tokenizer from '{MODEL_PATH}'...")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+    
+    print(f"Manually rebuilding MoE model structure...")
+    # First, we create an "empty shell" of our model with the correct architecture but random weights.
+    # `from_config` builds the structure without loading any weights.
+    # 首先，我们用`from_config`创建一个拥有正确架构但权重是随机的“空壳”模型。
+    # 这一步只搭建骨架，不加载任何权重。
+    moe_model = AutoModelForCausalLM.from_config(config)
+    
+    # Then, we perform the "architectural surgery" again, replacing standard MLPs with our MoEModules.
+    # 然后，我们再次手动进行“架构手术”，把标准的MLP替换成我们的MoE模块。
+    for i, layer in enumerate(moe_model.model.layers):
+        layer.mlp = MoEModule(config)
+        
+    print(f"Loading your trained MoE weights into the correct structure...")
+    from safetensors.torch import load_file
+    state_dict = load_file(os.path.join(MODEL_PATH, "model.safetensors"), device="cpu")
+    
+    # Use `strict=False` for flexible loading, then manually tie the weights.
+    # This handles the missing `lm_head.weight` key caused by weight tying.
+    # 使用`strict=False`进行灵活加载，然后手动绑定权重。
+    # 这个操作处理了因权重绑定而导致的`lm_head.weight`键缺失的问题。
+    moe_model.load_state_dict(state_dict, strict=False)
+    moe_model.tie_weights()
+    
+    # Move the finalized model to the GPU and set it to evaluation mode.
+    # 将最终完成的模型移动到GPU，并设置为评估模式。
+    moe_model.to(device, dtype=torch.bfloat16)
+    moe_model.eval()
+    print("--- MoE Model is ready for conversation! ---")
+    print("Type 'exit' or 'quit' to end the chat.\n")
+
+    # --- Interactive Conversation Loop / 交互式对话循环 ---
+    messages = []
+    while True:
+        try:
+            user_input = input("You: ")
+            if user_input.lower() in ["exit", "quit"]:
+                print("Goodbye!")
+                break
+
+            # Step 1: Add the user's input to the conversation history.
+            # 步骤 1: 将用户的输入添加到对话历史中。
+            messages.append({"role": "user", "content": user_input})
+            
+            # Step 2: Format the entire conversation history using the chat template.
+            # `add_generation_prompt=True` adds the starting tokens for the assistant's turn.
+            # 步骤 2: 使用聊天模板格式化完整的对话历史。
+            # `add_generation_prompt=True` 会在末尾添加助手角色的起始标记。
+            prompt_text = tokenizer.apply_chat_template(
+                messages, 
+                tokenize=False, 
+                add_generation_prompt=True
+            )
+            
+            # Step 3: Encode the input text and move it to the GPU.
+            # 步骤 3: 编码输入文本并将其发送到GPU。
+            inputs = tokenizer(prompt_text, return_tensors="pt").to(device)
+
+            # Step 4: Generate a response.
+            # 步骤 4: 生成回复。
+            with torch.no_grad():
+                outputs = moe_model.generate(
+                    **inputs, 
+                    max_new_tokens=256, 
+                    temperature=0.7, 
+                    top_p=0.9, 
+                    do_sample=True
+                )
+            
+            # Step 5: Decode and clean the output.
+            # `outputs[0]` contains the full conversation (input + output). We need to extract only the new part.
+            # 步骤 5: 解码并清理输出。
+            # `outputs[0]` 包含了完整的对话（输入+输出），我们需要从中提取出模型新生成的部分。
+            full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+            # Find the newly generated part by removing the original prompt from the full response.
+            # 通过从完整回复中移除原始的prompt来找到新生成的部分。
+            assistant_prompt_start = "<|assistant|>\n"
+            # This is a robust way to find the start of the assistant's actual response
+            assistant_response_start_index = full_response.rfind(assistant_prompt_start)
+            if assistant_response_start_index != -1:
+                model_response = full_response[assistant_response_start_index + len(assistant_prompt_start):].strip()
+            else:
+                 # Fallback for simpler cases
+                model_response = full_response.replace(prompt_text.replace("<s>", "").replace("</s>", ""), "").strip()
+
+            print(f"MoE Model: {model_response}")
+            
+            # Step 6: Add the model's response to the history for multi-turn conversations.
+            # 步骤 6: 将模型的回复也添加到对话历史中，以便进行多轮对话。
+            messages.append({"role": "assistant", "content": model_response})
+
+        except KeyboardInterrupt:
+            print("\nGoodbye!")
+            break
+
+# Script entry point / 脚本入口
+if __name__ == "__main__":
+    main()

train/test_moe_model.py

@@ -0,0 +1,155 @@
+# ==============================================================================
+#                 Smol-MoE 8x135M - "The Mind-Reader" Test Script
+#                     (Final Version with Correct Loading Logic)
+# ==============================================================================
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+from transformers.models.llama.modeling_llama import LlamaMLP, LlamaAttention, LlamaRMSNorm, LlamaForCausalLM
+import numpy as np
+import os
+# --- 1. 关键：重新定义你的所有自定义模块 ---
+# 我们需要这些类的定义，来手动重建模型的正确结构
+
+MODEL_PATH = "./SmolMoE-8x135M-Instruct-v1-Trained"
+config = AutoConfig.from_pretrained(MODEL_PATH)
+NUM_EXPERTS = config.moe_num_experts
+TOP_K = config.moe_top_k
+
+EXPERT_NAMES = [
+    "Actor", "Analyst", "Coder", "Encyclopedia",
+    "Guardian", "Summarizer", "Thinker", "Writer"
+]
+
+class MoERouter(nn.Module):
+    def __init__(self, hidden_size: int, num_experts: int):
+        super().__init__()
+        self.layer = nn.Linear(hidden_size, num_experts, bias=False)
+    def forward(self, hidden_states):
+        return self.layer(hidden_states)
+
+class MoEModule(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.top_k = TOP_K
+        self.num_experts = NUM_EXPERTS
+        self.router = MoERouter(self.hidden_size, self.num_experts)
+        self.experts = nn.ModuleList([LlamaMLP(config) for _ in range(self.num_experts)])
+
+    def forward(self, hidden_states):
+        original_shape = hidden_states.shape
+        flat_hidden_states = hidden_states.view(-1, self.hidden_size)
+        router_logits = self.router(flat_hidden_states)
+        routing_weights = F.softmax(router_logits, dim=-1, dtype=torch.float)
+        routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
+        routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
+        routing_weights = routing_weights.to(hidden_states.dtype)
+        final_hidden_states = torch.zeros_like(flat_hidden_states)
+        for k in range(self.top_k):
+            expert_indices_k = selected_experts[:, k]
+            routing_weights_k = routing_weights[:, k]
+            for i in range(self.num_experts):
+                mask = expert_indices_k == i
+                if mask.any():
+                    expert_output = self.experts[i](flat_hidden_states[mask])
+                    final_hidden_states.index_add_(0, torch.where(mask)[0], expert_output * routing_weights_k[mask].unsqueeze(1))
+        return final_hidden_states.view(*original_shape)
+
+# --- 2. 诊断测试的核心工具 ---
+captured_router_weights = {}
+def get_router_weights_hook(layer_idx):
+    """这是一个创建钩子函数的工厂"""
+    def hook(module, input, output):
+        # input[0] 是进入MoE模块的hidden_states
+        router_logits = module.router(input[0])
+        # 我们计算整个句子（所有token）的平均路由概率
+        avg_probs = F.softmax(router_logits, dim=-1).mean(dim=[0, 1])
+
+        # *** 这是最终的修复：在转换成numpy前，先转换成float32 ***
+        captured_router_weights[layer_idx] = avg_probs.detach().cpu().to(torch.float32).numpy()
+    return hook
+
+def visualize_router_decisions(prompt):
+    print("\n" + "="*80)
+    print(f"ROUTER DECISION ANALYSIS for Prompt: '{prompt[:50]}...'")
+    print("="*80)
+    print(f"{'Layer':<7} | {'Dominant Expert(s)':<45} | {'Confidence'}")
+    print("-"*80)
+    for layer_idx, weights in captured_router_weights.items():
+        top2_indices = np.argsort(weights)[-2:][::-1]
+        dominant_experts_str = f"1. {EXPERT_NAMES[top2_indices[0]]} | 2. {EXPERT_NAMES[top2_indices[1]]}"
+        confidence_str = f"({weights[top2_indices[0]]:.1%} | {weights[top2_indices[1]]:.1%})"
+        print(f"Layer {layer_idx:<4} | {dominant_experts_str:<45} | {confidence_str}")
+    print("="*80 + "\n")
+
+# --- 3. 主测试流程 ---
+def main():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    print(f"Loading tokenizer from '{MODEL_PATH}'...")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+    
+    print(f"Manually rebuilding MoE model structure...")
+    # 首先，我们用`from_config`创建一个随机权重的、但结构正确的“空壳”
+    moe_model = AutoModelForCausalLM.from_config(config)
+    
+    # 然后，我们手动进行“架构手术”，把标准的MLP替换成我们的MoE模块
+    for i, layer in enumerate(moe_model.model.layers):
+        layer.mlp = MoEModule(config)
+        
+    print(f"Loading your trained MoE weights into the correct structure...")
+    # 从safetensors文件加载权重
+    from safetensors.torch import load_file
+    state_dict = load_file(os.path.join(MODEL_PATH, "model.safetensors"), device="cpu")
+    
+    # *** 这是最终的修复：第一步 - 使用strict=False灵活加载 ***
+    # 我们知道lm_head.weight会缺失，所以允许这种“不严格”的加载
+    moe_model.load_state_dict(state_dict, strict=False)
+    
+    # *** 这是最终的修复：第二步 - 手动执行权重绑定 ***
+    # 这个函数会根据config中的"tie_word_embeddings"设置，将lm_head和词嵌入层绑定
+    moe_model.tie_weights()
+    
+    moe_model.to(device, dtype=torch.bfloat16)
+    moe_model.eval()
+    print("--- Custom MoE Model Successfully Loaded and Finalized! ---")
+
+    # 为诊断测试安装“窃听器”
+    hooks = []
+    for i, layer in enumerate(moe_model.model.layers):
+        # 确保我们是在MoEModule上注册钩子，而不是标准的LlamaMLP
+        if isinstance(layer.mlp, MoEModule):
+            hook = layer.mlp.register_forward_hook(get_router_weights_hook(i))
+            hooks.append(hook)
+
+    # 设计一系列“考题”
+    test_prompts = {
+        "Coder": "Write a Python function that takes a list of numbers and returns a new list with only the even numbers.",
+        "Writer": "In a world where shadows have a life of their own, a young lamplighter discovers a terrible secret. Write the opening paragraph.",
+        "Thinker": "If all bloops are gloops, and some gloops are zloops, is it certain that some bloops are zloops? Explain your reasoning.",
+        "Encyclopedia": "What were the primary economic and political causes of the French Revolution?",
+        "Multi-Expert": "In the style of a Shakespearean tragedy, write a short monologue for a software developer lamenting a bug in their code. Include a comment line from the code."
+    }
+
+    for expert_name, prompt in test_prompts.items():
+        captured_router_weights.clear()
+        print(f"\n--- Testing for: {expert_name} Expert ---")
+        print(f"Prompt: {prompt}")
+        inputs = tokenizer(prompt, return_tensors="pt").to(device)
+        with torch.no_grad():
+            outputs = moe_model.generate(**inputs, max_new_tokens=100, do_sample=True, top_k=50, top_p=0.95)
+        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        print("\n--- Generated Text ---")
+        print(generated_text)
+        print("--- End of Generated Text ---")
+        visualize_router_decisions(prompt)
+    
+    for hook in hooks:
+        hook.remove()
+    print("All tests complete and hooks have been removed.")
+
+if __name__ == "__main__":
+    main()

train/test_moe_model_zhcn_en.py

@@ -0,0 +1,205 @@
+# ==============================================================================
+#                 Smol-MoE 8x135M - "The Mind-Reader" Test Script
+#                     (Final Version with Correct Loading Logic)
+#
+#                 Smol-MoE 8x135M - “读心器”测试脚本
+#                     (包含正确加载逻辑的最终版本)
+# ==============================================================================
+
+# --- Core Library Imports / 核心库导入 ---
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+from transformers.models.llama.modeling_llama import LlamaMLP, LlamaForCausalLM # We need LlamaMLP for the expert definition / 我们需要 LlamaMLP 来定义专家
+import numpy as np
+import os # We need this for path operations / 我们需要 os 库来处理文件路径
+
+# --- 1. CRITICAL: Re-define Your Custom Architecture ---
+# --- 1. 关键：重新定义你的所有自定义模块 ---
+# When loading a model with a custom architecture, Hugging Face needs to know the definition of the custom classes.
+# By defining them here, we allow `from_pretrained` to correctly reconstruct our unique MoE model.
+# 在加载一个拥有自定义架构的模型时，Hugging Face 需要知道这些自定义类的定义。
+# 在这里定义它们，我们才能让 `from_pretrained` 函数成功地重建我们独一无二的MoE模型。
+
+# --- Model Configuration / 模型配置 ---
+MODEL_PATH = "./SmolMoE-8x135M-Instruct-v1-Trained"
+config = AutoConfig.from_pretrained(MODEL_PATH)
+# Load our custom MoE parameters from the saved config file.
+# 从我们保存的配置文件中，加载自定义的MoE参数。
+NUM_EXPERTS = config.moe_num_experts
+TOP_K = config.moe_top_k
+
+# A list of expert names for clear visualization later. The order must match the training script.
+# 用于后续清晰可视化的专家名称列表。顺序必须和训练脚本中的保持一致。
+EXPERT_NAMES = [
+    "Actor", "Analyst", "Coder", "Encyclopedia",
+    "Guardian", "Summarizer", "Thinker", "Writer"
+]
+
+class MoERouter(nn.Module):
+    """The Router module. Its job is to score experts for each token."""
+    """路由器模块。它的工作是为每个token给所有专家打分。"""
+    def __init__(self, hidden_size: int, num_experts: int):
+        super().__init__()
+        self.layer = nn.Linear(hidden_size, num_experts, bias=False)
+    def forward(self, hidden_states):
+        return self.layer(hidden_states)
+
+class MoEModule(nn.Module):
+    """The custom Mixture-of-Experts module that replaces the standard FFN."""
+    """我们自定义的混合专家模块，它替换了标准的FFN。"""
+    def __init__(self, config):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.top_k = TOP_K
+        self.num_experts = NUM_EXPERTS
+        self.router = MoERouter(self.hidden_size, self.num_experts)
+        self.experts = nn.ModuleList([LlamaMLP(config) for _ in range(self.num_experts)])
+
+    def forward(self, hidden_states):
+        original_shape = hidden_states.shape
+        flat_hidden_states = hidden_states.view(-1, self.hidden_size)
+        router_logits = self.router(flat_hidden_states)
+        routing_weights = F.softmax(router_logits, dim=-1, dtype=torch.float)
+        routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
+        routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
+        routing_weights = routing_weights.to(hidden_states.dtype)
+        final_hidden_states = torch.zeros_like(flat_hidden_states)
+        for k in range(self.top_k):
+            expert_indices_k = selected_experts[:, k]
+            routing_weights_k = routing_weights[:, k]
+            for i in range(self.num_experts):
+                mask = expert_indices_k == i
+                if mask.any():
+                    expert_output = self.experts[i](flat_hidden_states[mask])
+                    final_hidden_states.index_add_(0, torch.where(mask)[0], expert_output * routing_weights_k[mask].unsqueeze(1))
+        return final_hidden_states.view(*original_shape)
+
+# --- 2. Core Diagnostic Tools / 诊断测试的核心工具 ---
+# This dictionary will store the router decisions captured by our hooks.
+# 这个字典将用于存储我们的钩子捕获到的路由器决策数据。
+captured_router_weights = {}
+
+def get_router_weights_hook(layer_idx):
+    """This is a factory function that creates our hook."""
+    """这是一个创建钩子函数的工厂函数。"""
+    def hook(module, input, output):
+        # `input[0]` is the hidden_states tensor passed to the MoE module.
+        # `input[0]` 是传入MoE模块的hidden_states张量。
+        router_logits = module.router(input[0])
+        # We calculate the average routing probability for all tokens in the sequence.
+        # 我们计算序列中所有token的平��路由概率。
+        avg_probs = F.softmax(router_logits, dim=-1).mean(dim=[0, 1])
+
+        # *** FINAL FIX: Convert from BFloat16 to Float32 before converting to NumPy. ***
+        # NumPy does not support the bfloat16 dtype, so we must convert it first.
+        # *** 最终修复：在转换为numpy数组前，先将数据格式从BFloat16转换为Float32。***
+        # NumPy库不支持bfloat16这种数据类型，所以我们必须先进行转换。
+        captured_router_weights[layer_idx] = avg_probs.detach().cpu().to(torch.float32).numpy()
+    return hook
+
+def visualize_router_decisions(prompt):
+    """A helper function to print the captured router decisions in a nice table."""
+    """一个辅助函数，用于将捕获到的路由器决策以漂亮的表格形式打印出来。"""
+    print("\n" + "="*80)
+    print(f"ROUTER DECISION ANALYSIS for Prompt: '{prompt[:50]}...'")
+    print("="*80)
+    print(f"{'Layer':<7} | {'Dominant Expert(s)':<45} | {'Confidence'}")
+    print("-"*80)
+    for layer_idx, weights in captured_router_weights.items():
+        top2_indices = np.argsort(weights)[-2:][::-1]
+        dominant_experts_str = f"1. {EXPERT_NAMES[top2_indices[0]]} | 2. {EXPERT_NAMES[top2_indices[1]]}"
+        confidence_str = f"({weights[top2_indices[0]]:.1%} | {weights[top2_indices[1]]:.1%})"
+        print(f"Layer {layer_idx:<4} | {dominant_experts_str:<45} | {confidence_str}")
+    print("="*80 + "\n")
+
+# --- 3. Main Testing Workflow / 主测试流程 ---
+def main():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    print(f"Loading tokenizer from '{MODEL_PATH}'...")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+    
+    print(f"Manually rebuilding MoE model structure...")
+    # First, we create an "empty shell" of our model with the correct architecture but random weights.
+    # `from_config` builds the structure without loading any weights.
+    # 首先，我们用`from_config`创建一个拥有正确架构但权重是随机的“空壳”模型。
+    # 这一步只搭建骨架，不加载任何权重。
+    moe_model = AutoModelForCausalLM.from_config(config)
+    
+    # Then, we perform the "architectural surgery" again, replacing standard MLPs with our MoEModules.
+    # 然后，我们再次手动进行“架构手术”，把标准的MLP替换成我们的MoE模块。
+    for i, layer in enumerate(moe_model.model.layers):
+        layer.mlp = MoEModule(config)
+        
+    print(f"Loading your trained MoE weights into the correct structure...")
+    # Load the weights from the safetensors file.
+    # 从safetensors文件加载权重。
+    from safetensors.torch import load_file
+    state_dict = load_file(os.path.join(MODEL_PATH, "model.safetensors"), device="cpu")
+    
+    # *** FINAL FIX #1: Use `strict=False` for flexible loading. ***
+    # We know `lm_head.weight` is missing because of `tie_word_embeddings`, so we allow this "inexact" loading.
+    # *** 最终修复 #1：使用`strict=False`进行灵活加载。***
+    # 我们知道因为`tie_word_embeddings`的设置，`lm_head.weight`是缺失的，所以我们允许这种“不严格”的加载。
+    moe_model.load_state_dict(state_dict, strict=False)
+    
+    # *** FINAL FIX #2: Manually tie the weights. ***
+    # This function reads the `tie_word_embeddings` setting from the config and correctly links the lm_head to the token embeddings.
+    # *** 最终修复 #2：手动执行权重绑定。***
+    # 这个函数会根据config中的`tie_word_embeddings`设置，将lm_head和词嵌入层正确地绑定在一起。
+    moe_model.tie_weights()
+    
+    # Move the finalized model to the GPU and set it to evaluation mode.
+    # 将最终完成的模型移动到GPU，并设置为评估模式。
+    moe_model.to(device, dtype=torch.bfloat16)
+    moe_model.eval()
+    print("--- Custom MoE Model Successfully Loaded and Finalized! ---")
+
+    # Install our "listening devices" (hooks) on each MoE layer for diagnostics.
+    # 为诊断测试，在每个MoE层上都安装我们的“窃听器”（钩子）。
+    hooks = []
+    for i, layer in enumerate(moe_model.model.layers):
+        if isinstance(layer.mlp, MoEModule):
+            hook = layer.mlp.register_forward_hook(get_router_weights_hook(i))
+            hooks.append(hook)
+
+    # Design a series of "exam questions" to test different experts.
+    # 设计一系列“考题”来测试不同的专家。
+    test_prompts = {
+        "Coder": "Write a Python function that takes a list of numbers and returns a new list with only the even numbers.",
+        "Writer": "In a world where shadows have a life of their own, a young lamplighter discovers a terrible secret. Write the opening paragraph.",
+        "Thinker": "If all bloops are gloops, and some gloops are zloops, is it certain that some bloops are zloops? Explain your reasoning.",
+        "Encyclopedia": "What were the primary economic and political causes of the French Revolution?",
+        "Multi-Expert": "In the style of a Shakespearean tragedy, write a short monologue for a software developer lamenting a bug in their code. Include a comment line from the code."
+    }
+
+    # The main testing loop.
+    # 主测试循环。
+    for expert_name, prompt in test_prompts.items():
+        captured_router_weights.clear() # Clear data from the previous run / 清空上一次的捕获数据
+        print(f"\n--- Testing for: {expert_name} Expert ---")
+        print(f"Prompt: {prompt}")
+        inputs = tokenizer(prompt, return_tensors="pt").to(device)
+        
+        # 1. Functional Test: Generate text / 1. 功能测试：生成文本
+        with torch.no_grad():
+            outputs = moe_model.generate(**inputs, max_new_tokens=100, do_sample=True, top_k=50, top_p=0.95)
+        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        print("\n--- Generated Text ---")
+        print(generated_text)
+        print("--- End of Generated Text ---")
+        
+        # 2. Diagnostic Test: Visualize router decisions / 2. 诊断测试：可视化路由决策
+        visualize_router_decisions(prompt)
+    
+    # Clean up by removing all hooks to prevent memory leaks.
+    # 清理工作：移除所有钩子以防止内存泄漏。
+    for hook in hooks:
+        hook.remove()
+    print("All tests complete and hooks have been removed.")
+
+# Script entry point / 脚本入口
+if __name__ == "__main__":
+    main()

train/train_moe_router.py

@@ -0,0 +1,163 @@
+# ==============================================================================
+#                      Smol-MoE 8x135M - Master Script
+#                 (Final Version, All Fixes Included)
+# ==============================================================================
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.nn.functional as F
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer # <<< 这是最关键的修复！确保AutoTokenizer在这里
+from transformers.models.llama.modeling_llama import LlamaMLP
+from safetensors.torch import load_file # <<< 使用正确的safetensors加载器
+import os
+import shutil
+import numpy as np
+import time
+
+# --- 0. Configuration & Setup ---
+MODEL_NAME = "./SmolLM2-135M-Instruct"
+BASE_EXPERT_PATH = "./models"
+EXPERT_DIRS = [
+    "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"
+]
+NUM_EXPERTS = 8
+TOP_K = 2
+LEARNING_RATE = 0.001
+EPOCHS = 20 # 既然我们知道模拟数据无法让模型学习，20轮足以验证流程
+BATCH_SIZE = 4
+SEQUENCE_LENGTH = 128
+LB_LOSS_COEFFICIENT = 0.01
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
+
+
+# --- 1. Define the MoE Architecture Components ---
+class MoERouter(nn.Module):
+    def __init__(self, hidden_size: int, num_experts: int):
+        super().__init__()
+        self.layer = nn.Linear(hidden_size, num_experts, bias=False)
+    def forward(self, hidden_states):
+        return self.layer(hidden_states)
+
+class MoEModule(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.top_k = TOP_K
+        self.num_experts = NUM_EXPERTS
+        self.router = MoERouter(self.hidden_size, self.num_experts)
+        self.experts = nn.ModuleList([LlamaMLP(config) for _ in range(self.num_experts)])
+        self.most_recent_lb_loss = None
+
+    def forward(self, hidden_states):
+        original_shape = hidden_states.shape
+        flat_hidden_states = hidden_states.view(-1, self.hidden_size)
+        router_logits = self.router(flat_hidden_states)
+        routing_weights = F.softmax(router_logits, dim=-1, dtype=torch.float)
+        routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
+        routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
+        routing_weights = routing_weights.to(hidden_states.dtype)
+        router_probs_full = F.softmax(router_logits, dim=-1, dtype=torch.float)
+        avg_expert_prob = router_probs_full.mean(dim=0)
+        expert_mask_for_lb = F.one_hot(selected_experts, num_classes=self.num_experts).sum(dim=1)
+        avg_expert_fraction = expert_mask_for_lb.float().mean(dim=0)
+        self.most_recent_lb_loss = self.num_experts * torch.sum(avg_expert_prob * avg_expert_fraction)
+        final_hidden_states = torch.zeros_like(flat_hidden_states)
+        for k in range(self.top_k):
+            expert_indices_k = selected_experts[:, k]
+            routing_weights_k = routing_weights[:, k]
+            for i in range(self.num_experts):
+                mask = expert_indices_k == i
+                if mask.any():
+                    expert_output = self.experts[i](flat_hidden_states[mask])
+                    final_hidden_states.index_add_(0, torch.where(mask)[0], expert_output * routing_weights_k[mask].unsqueeze(1))
+        return final_hidden_states.view(*original_shape)
+
+# --- 2. The Grand Assembly Function ---
+def create_moe_model():
+    print("--- Starting Architectural Surgery ---")
+    config = AutoConfig.from_pretrained(MODEL_NAME)
+    print("Step 1: Loading base model skeleton...")
+    base_model = AutoModelForCausalLM.from_pretrained(
+        os.path.join(BASE_EXPERT_PATH, EXPERT_DIRS[0]),
+        torch_dtype=torch.bfloat16,
+        device_map=device
+    )
+    print("Step 2: Pre-loading all expert weights into CPU memory for efficiency...")
+    all_experts_state_dicts = [
+        load_file(os.path.join(BASE_EXPERT_PATH, expert_dir, 'model.safetensors'), device='cpu')
+        for expert_dir in EXPERT_DIRS
+    ]
+    print("All expert weights pre-loaded.")
+    print("Step 3: Replacing FFNs with MoE modules and transplanting expert weights...")
+    for layer_idx, layer in enumerate(base_model.model.layers):
+        layer.mlp = MoEModule(config).to(device, dtype=torch.bfloat16)
+        for expert_idx in range(NUM_EXPERTS):
+            expert_state_dict = all_experts_state_dicts[expert_idx]
+            expert_mlp_weights = {
+                k.replace(f"model.layers.{layer_idx}.mlp.", ""): v
+                for k, v in expert_state_dict.items()
+                if f"model.layers.{layer_idx}.mlp." in k
+            }
+            layer.mlp.experts[expert_idx].load_state_dict(expert_mlp_weights)
+    print("Step 4: Freezing all parameters except for the routers...")
+    for name, param in base_model.named_parameters():
+        if "router" not in name:
+            param.requires_grad = False
+    print("\n--- Surgery Complete! MoE Model is assembled and ready for training. ---")
+    trainable_params = sum(p.numel() for p in base_model.parameters() if p.requires_grad)
+    total_params = sum(p.numel() for p in base_model.parameters())
+    print(f"Total Parameters: {total_params / 1e6:.2f}M")
+    print(f"Trainable Parameters (Routers): {trainable_params}")
+    return base_model
+
+# --- 3. The Main Training & Saving Process ---
+def main():
+    moe_model = create_moe_model()
+    optimizer = optim.AdamW([p for p in moe_model.parameters() if p.requires_grad], lr=LEARNING_RATE)
+    print("\n--- Preparing Simulated Mixed Dataset for Training ---")
+    mock_input_ids = torch.randint(0, moe_model.config.vocab_size, (BATCH_SIZE, SEQUENCE_LENGTH), device=device)
+    mock_labels = mock_input_ids.clone()
+    print("--- Starting Router Training Loop (Optimized & Corrected) ---")
+    moe_model.train()
+    start_time = time.time()
+    for epoch in range(EPOCHS):
+        optimizer.zero_grad()
+        outputs = moe_model(input_ids=mock_input_ids, labels=mock_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()
+        if (epoch + 1) % 10 == 0:
+            elapsed_time = time.time() - start_time
+            print(f"Epoch [{epoch+1:03d}/{EPOCHS}] | 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: {elapsed_time:.2f}s")
+            start_time = time.time()
+    print("\n--- Router Training Complete! ---")
+    print("\n--- Phase 5: Saving the fully trained MoE model to disk ---")
+    OUTPUT_MODEL_DIR = "./SmolMoE-8x135M-Instruct-v1-Trained"
+    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 = AutoTokenizer.from_pretrained(MODEL_NAME)
+    tokenizer.save_pretrained(OUTPUT_MODEL_DIR)
+    print("\n--- Model successfully saved! ---")
+    print("You can now load this model in other scripts, but you must re-define the custom MoE classes first.")
+
+if __name__ == "__main__":
+    main()

train/train_moe_router_en.py

@@ -0,0 +1,289 @@
+# ==============================================================================
+#                      Smol-MoE 8x135M - The "Genesis" Master Script
+#                 (Final Optimized Version with All Fixes & Detailed Comments)
+# ==============================================================================
+
+# --- Core Library Imports ---
+# PyTorch Core
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.nn.functional as F
+
+# Hugging Face Transformers library, the source of our "Lego bricks"
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+from transformers.models.llama.modeling_llama import LlamaMLP # The standard FFN module from Llama, which we use as the base for our "Experts"
+
+# Safetensors library, for safely and efficiently loading model weights
+from safetensors.torch import load_file
+
+# Standard Python Libraries
+import os       # For handling file paths
+import shutil   # For directory operations (like deleting an old model folder)
+import numpy as np    # For data analysis (in the final test)
+import time     # For timing the training process
+
+# --- 0. Global Configuration & Hyperparameters ---
+# This is the master control panel for the entire project. All key parameters are defined here.
+
+# MODEL_NAME: The path to the base model. We use this to load the initial config and tokenizer.
+# Note: This should point to a standard, unmodified SmolLM model.
+MODEL_NAME = "./SmolLM2-135M-Instruct"
+
+# BASE_EXPERT_PATH: The parent directory containing all 8 of your pre-trained expert model folders.
+BASE_EXPERT_PATH = "./models"
+
+# EXPERT_DIRS: A list of the specific directory names for your 8 expert models. The order is important.
+EXPERT_DIRS = [
+    "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"
+]
+
+# MoE Architecture Parameters
+NUM_EXPERTS = 8  # The number of experts in our committee
+TOP_K = 2        # The number of top experts to route to for each token
+
+# Training Hyperparameters
+LEARNING_RATE = 0.001       # The learning rate for the routers. Since we only train routers, it can be slightly higher.
+EPOCHS = 20                 # Number of training epochs. Since we're using mock data to validate the process, 20 is sufficient.
+                            # This needs to be much higher when using real data.
+BATCH_SIZE = 4              # The number of sequences to process in each batch. Adjust based on your VRAM.
+SEQUENCE_LENGTH = 128       # The length of text sequences the model processes. Adjust based on your VRAM.
+LB_LOSS_COEFFICIENT = 0.01  # The weight coefficient for the load balancing loss. This is a critical "balancing valve"
+                            # used to trade off between "doing the job well" and "distributing work fairly."
+
+# Device Configuration
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
+
+
+# --- 1. MoE Architecture Component Definitions ---
+# These are the blueprints for the new parts we've "invented."
+
+class MoERouter(nn.Module):
+    """
+    The Router (or Gate Network) - The "CEO" or "dispatcher" of the expert committee.
+    Its structure is a simple linear layer, responsible for scoring all experts for each incoming token.
+    """
+    def __init__(self, hidden_size: int, num_experts: int):
+        super().__init__()
+        self.layer = nn.Linear(hidden_size, num_experts, bias=False)
+
+    def forward(self, hidden_states):
+        # Outputs the "scores" (logits) for each expert, which will later be turned into probabilities via Softmax.
+        return self.layer(hidden_states)
+
+class MoEModule(nn.Module):
+    """
+    The Mixture-of-Experts Module - The "conference room" for the entire expert committee.
+    This module replaces the standard FFN (MLP) block in the original Llama model.
+    It contains one router (the CEO) and a list of experts (the board members).
+    """
+    def __init__(self, config):
+        super().__init__()
+        # Get necessary parameters from the global config
+        self.hidden_size = config.hidden_size
+        self.top_k = TOP_K
+        self.num_experts = NUM_EXPERTS
+        
+        # Create the components
+        self.router = MoERouter(self.hidden_size, self.num_experts)
+        # LlamaMLP is the standard FFN implementation in Hugging Face's Llama, which we use as the base for our "experts".
+        self.experts = nn.ModuleList([LlamaMLP(config) for _ in range(self.num_experts)])
+        
+        # A placeholder to temporarily store the load balancing loss for this layer during a forward pass
+        self.most_recent_lb_loss = None
+
+    def forward(self, hidden_states):
+        # Store the original shape to reshape the output at the end
+        original_shape = hidden_states.shape
+        # Flatten the input from (batch, sequence, dim) to (batch * sequence, dim) for token-level routing
+        flat_hidden_states = hidden_states.view(-1, self.hidden_size)
+
+        # --- Step 1: Routing Decision ---
+        # Get scores from the router for each token
+        router_logits = self.router(flat_hidden_states)
+        # Use Softmax to convert scores to probabilities
+        routing_weights = F.softmax(router_logits, dim=-1, dtype=torch.float)
+        # Select the top-k experts and their corresponding probabilities
+        routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
+        # Normalize the probabilities of the top-k experts so they sum to 1
+        routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
+        # Cast weights back to the model's main dtype (e.g., bfloat16) for efficiency
+        routing_weights = routing_weights.to(hidden_states.dtype)
+
+        # --- Step 2: Calculate and Store Load Balancing Loss ---
+        # This is the soul of MoE training: ensuring the router doesn't get "lazy" and uses all experts fairly.
+        router_probs_full = F.softmax(router_logits, dim=-1, dtype=torch.float)
+        avg_expert_prob = router_probs_full.mean(dim=0) # The average probability for each expert across all tokens
+        expert_mask_for_lb = F.one_hot(selected_experts, num_classes=self.num_experts).sum(dim=1) # Checks which experts were chosen for each token
+        avg_expert_fraction = expert_mask_for_lb.float().mean(dim=0) # The average fraction of tokens processed by each expert
+        # Calculate the loss, multiply by the number of experts as a penalty term, and store it.
+        self.most_recent_lb_loss = self.num_experts * torch.sum(avg_expert_prob * avg_expert_fraction)
+
+        # --- Step 3: Expert Computation and Result Aggregation (Vectorized & Efficient) ---
+        # Create an empty tensor to store the final results
+        final_hidden_states = torch.zeros_like(flat_hidden_states)
+        
+        # This loop only iterates `top_k` times (e.g., 2), which is very fast.
+        for k in range(self.top_k):
+            # Get the expert indices and weights for the k-th choice across all tokens
+            expert_indices_k = selected_experts[:, k]
+            routing_weights_k = routing_weights[:, k]
+            
+            # This loop iterates over all experts, but the computations inside are batched and fast.
+            for i in range(self.num_experts):
+                # Create a mask to find all tokens that were routed to the current expert `i`
+                mask = expert_indices_k == i
+                if mask.any(): # If any token was routed to this expert
+                    # Process all selected tokens in a single batch
+                    expert_output = self.experts[i](flat_hidden_states[mask])
+                    # Weight the expert's output by its routing weight and "add" it back to the correct positions in the final tensor
+                    final_hidden_states.index_add_(0, torch.where(mask)[0], expert_output * routing_weights_k[mask].unsqueeze(1))
+        
+        # Reshape the result back to the original (batch, sequence, hidden_size) shape and return
+        return final_hidden_states.view(*original_shape)
+
+# --- 2. The "Genesis" Function: Assembling, Transplanting, and Modifying the Model ---
+def create_moe_model():
+    """
+    This is the "Architectural Surgery" function. It is responsible for:
+    1. Building an empty model skeleton with MoE modules.
+    2. "Transplanting" the weights from your 8 pre-trained experts into it.
+    3. Freezing all expert parameters, leaving only the routers trainable.
+    """
+    print("--- Starting Architectural Surgery ---")
+    
+    # Load the config from the standard model; this is the "genetic blueprint" for our new model
+    config = AutoConfig.from_pretrained(MODEL_NAME)
+    
+    print("Step 1: Loading base model skeleton...")
+    # Load one of the experts to serve as the "skeleton" for our MoE model.
+    # We will use its non-FFN parts (embeddings, attention modules, etc.).
+    base_model = AutoModelForCausalLM.from_pretrained(
+        os.path.join(BASE_EXPERT_PATH, EXPERT_DIRS[0]),
+        torch_dtype=torch.bfloat16,
+        device_map=device
+    )
+    
+    print("Step 2: Pre-loading all expert weights into CPU memory for efficiency...")
+    # To improve efficiency, we load all expert weights from disk into CPU RAM at once.
+    # We use `safetensors.torch.load_file` as it is the correct and safe way to load .safetensors files.
+    all_experts_state_dicts = [
+        load_file(os.path.join(BASE_EXPERT_PATH, expert_dir, 'model.safetensors'), device='cpu')
+        for expert_dir in EXPERT_DIRS
+    ]
+    print("All expert weights pre-loaded.")
+
+    print("Step 3: Replacing FFNs with MoE modules and transplanting expert weights...")
+    # Iterate through all 30 layers of the model
+    for layer_idx, layer in enumerate(base_model.model.layers):
+        # In each layer, replace the original, standard LlamaMLP with our custom MoEModule
+        layer.mlp = MoEModule(config).to(device, dtype=torch.bfloat16)
+        
+        # Begin the "Organ Transplant"
+        for expert_idx in range(NUM_EXPERTS):
+            # Get the weights for the current expert from memory
+            expert_state_dict = all_experts_state_dicts[expert_idx]
+            # Filter to get only the weights for the FFN part of the current layer
+            expert_mlp_weights = {
+                k.replace(f"model.layers.{layer_idx}.mlp.", ""): v
+                for k, v in expert_state_dict.items()
+                if f"model.layers.{layer_idx}.mlp." in k
+            }
+            # Load these weights into the corresponding expert "seat" in our MoE module
+            layer.mlp.experts[expert_idx].load_state_dict(expert_mlp_weights)
+            
+    print("Step 4: Freezing all parameters except for the routers...")
+    # This is our key strategy: only train the "CEO", don't disturb the already-smart "experts".
+    for name, param in base_model.named_parameters():
+        if "router" not in name:
+            param.requires_grad = False
+            
+    print("\n--- Surgery Complete! MoE Model is assembled and ready for training. ---")
+    # Print parameter statistics to verify our operation was successful
+    trainable_params = sum(p.numel() for p in base_model.parameters() if p.requires_grad)
+    total_params = sum(p.numel() for p in base_model.parameters())
+    print(f"Total Parameters: {total_params / 1e6:.2f}M")
+    print(f"Trainable Parameters (Routers): {trainable_params}")
+    
+    return base_model
+
+# --- 3. Main Process: Training and Saving ---
+def main():
+    # Step 1: Call the "Genesis" function to create our model
+    moe_model = create_moe_model()
+    
+    # Step 2: Create the optimizer. It's smart enough to only include parameters where `requires_grad=True` (i.e., the routers).
+    optimizer = optim.AdamW([p for p in moe_model.parameters() if p.requires_grad], lr=LEARNING_RATE)
+    
+    print("\n--- Preparing Simulated Mixed Dataset for Training ---")
+    # NOTE: We are using completely random "mock data" here, solely to validate that the entire process runs.
+    # To make the routers truly intelligent, you MUST replace this with a real, diverse dataset.
+    mock_input_ids = torch.randint(0, moe_model.config.vocab_size, (BATCH_SIZE, SEQUENCE_LENGTH), device=device)
+    mock_labels = mock_input_ids.clone()
+    
+    print("--- Starting Router Training Loop (Optimized & Corrected) ---")
+    moe_model.train() # Set the model to training mode
+    
+    start_time = time.time()
+    for epoch in range(EPOCHS):
+        optimizer.zero_grad() # Clear gradients from the previous epoch
+        
+        # --- The Elegant and Correct Forward Pass ---
+        # Call the model directly. Hugging Face automatically handles all complex internal details (like attention masks).
+        # By providing `labels`, it also automatically calculates the main cross-entropy loss for us.
+        outputs = moe_model(input_ids=mock_input_ids, labels=mock_labels)
+        main_loss = outputs.loss # Extract the main task loss
+        
+        # --- Safely Collect Load Balancing Losses ---
+        total_lb_loss = 0.0
+        for layer in moe_model.model.layers:
+            total_lb_loss += layer.mlp.most_recent_lb_loss # Retrieve the loss stored in our placeholder
+            
+        # --- Calculate the Final "Composite KPI" (Total Loss) ---
+        total_loss = main_loss + LB_LOSS_COEFFICIENT * total_lb_loss
+        
+        # --- Backpropagation and Optimization ---
+        total_loss.backward() # Calculate gradients
+        optimizer.step()      # Update router weights
+        
+        # --- Print Training Logs ---
+        if (epoch + 1) % 10 == 0:
+            elapsed_time = time.time() - start_time
+            print(f"Epoch [{epoch+1:03d}/{EPOCHS}] | 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: {elapsed_time:.2f}s")
+            start_time = time.time()
+            
+    print("\n--- Router Training Complete! ---")
+    
+    # --- Step 5: Solidifying our great work onto the disk ---
+    print("\n--- Phase 5: Saving the fully trained MoE model to disk ---")
+    OUTPUT_MODEL_DIR = "./SmolMoE-8x135M-Instruct-v1-Trained"
+    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...")
+    # We use custom names for our MoE parameters to avoid conflicts with standard Hugging Face generation configs.
+    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) # Saves weights and the updated config file
+    
+    print("Saving tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) # Save the tokenizer
+    tokenizer.save_pretrained(OUTPUT_MODEL_DIR)
+    
+    print("\n--- Model successfully saved! ---")
+    print("You can now load this model in other scripts, but you must re-define the custom MoE classes first.")
+
+# --- Script Entry Point ---
+# Ensures that the main() function is only called when this file is executed directly.
+if __name__ == "__main__":
+    main()

train/train_moe_router_zh_CN.py

@@ -0,0 +1,289 @@
+# ==============================================================================
+#                      Smol-MoE 8x135M - “创世纪”主脚本
+#                 (最终优化版，包含所有修正与详细注释)
+# ==============================================================================
+
+# --- 核心库导入 ---
+# PyTorch 核心
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.nn.functional as F
+
+# Hugging Face Transformers 库，我们的“乐高积木”来源
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+from transformers.models.llama.modeling_llama import LlamaMLP # Llama模型中标准的FFN模块，我们将其作为“专家”的基础结构
+
+# Safetensors 库，用于安全、高效地加载模型权重
+from safetensors.torch import load_file
+
+# 标准Python库
+import os       # 用于处理文件路径
+import shutil   # 用于处理文件夹（如删除旧的模型文件夹）
+import numpy as np    # 用于数据分析（在最终测试中）
+import time     # 用于计算训练时长
+
+# --- 0. 全局配置与超参数 ---
+# 这里是整个项目的“总控制台”，所有关键参数都在此定义。
+
+# MODEL_NAME: 基础模型的路径。我们用它来加载初始的配置(config)和分词器(tokenizer)。
+# 注意：它指向一个标准的、未经修改的SmolLM模型。
+MODEL_NAME = "./SmolLM2-135M-Instruct"
+
+# BASE_EXPERT_PATH: 存放你所有8个预训练好的专家模型的父文件夹。
+BASE_EXPERT_PATH = "./models"
+
+# EXPERT_DIRS: 8个专家模型文件夹的具体名称列表。顺序很重要。
+EXPERT_DIRS = [
+    "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"
+]
+
+# MoE 架构参数
+NUM_EXPERTS = 8  # 专家委员会的专家数量
+TOP_K = 2        # 每次路由时，为每个Token选择的最优专家的数量
+
+# 训练超参数
+LEARNING_RATE = 0.001       # 路由器的学习率。只训练路由器，所以可以设置得稍高一些。
+EPOCHS = 20                 # 训练轮次。因为我们用的是模拟数据来验证流程，所以20轮就足够了。
+                            # 当使用真实数据时，你需要把它设置得更高。
+BATCH_SIZE = 4              # 每批次处理的数据量。根据你的显存大小调整。
+SEQUENCE_LENGTH = 128       # 模型处理的文本序列长度。根据你的显存大小调整。
+LB_LOSS_COEFFICIENT = 0.01  # 负载均衡损失的权重系数。这是个关键的“平衡阀”，
+                            # 用来平衡“任务做得好”和“分配得公平”这两个目标。
+
+# 设备配置
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
+
+
+# --- 1. MoE架构组件定义 ---
+# 这里是我们“发明”的新零件的蓝图。
+
+class MoERouter(nn.Module):
+    """
+    路由器 (Router / Gate Network) - 专家委员会的“CEO”或“调度员”。
+    它的结构非常简单，就是一个线性层，负责为每个输入的Token，给所有专家打分。
+    """
+    def __init__(self, hidden_size: int, num_experts: int):
+        super().__init__()
+        self.layer = nn.Linear(hidden_size, num_experts, bias=False)
+
+    def forward(self, hidden_states):
+        # 输出每个专家的“得分”(logits)，后续会通过Softmax转换成概率
+        return self.layer(hidden_states)
+
+class MoEModule(nn.Module):
+    """
+    混合专家模块 (MoE Module) - 整个“专家委员会”的会议室。
+    它替换了原来Llama模型中标准的FFN(MLP)模块。
+    它内部包含一个路由器(CEO)和8个专家(董事会成员)。
+    """
+    def __init__(self, config):
+        super().__init__()
+        # 从全局配置中获取必要的参数
+        self.hidden_size = config.hidden_size
+        self.top_k = TOP_K
+        self.num_experts = NUM_EXPERTS
+        
+        # 创建组件
+        self.router = MoERouter(self.hidden_size, self.num_experts)
+        # LlamaMLP是Hugging Face中Llama模型标准的FFN实现，我们用它作为“专家”的基础结构
+        self.experts = nn.ModuleList([LlamaMLP(config) for _ in range(self.num_experts)])
+        
+        # 创建一个占位符，用于在训练时临时存储该层的负载均衡损失
+        self.most_recent_lb_loss = None
+
+    def forward(self, hidden_states):
+        # 记录输入的原始形状，以便最后恢复
+        original_shape = hidden_states.shape
+        # 将输入“扁平化”处理，把batch和sequence维度合并，方便进行Token级别的路由
+        flat_hidden_states = hidden_states.view(-1, self.hidden_size)
+
+        # --- 步骤 1: 路由决策 ---
+        # 让路由器为每个Token打分
+        router_logits = self.router(flat_hidden_states)
+        # 用Softmax将分数转换成概率
+        routing_weights = F.softmax(router_logits, dim=-1, dtype=torch.float)
+        # 选出得分最高的Top-K个专家及其对应的概率
+        routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
+        # 对Top-K个专家的概率进行归一化，确保它们的和为1
+        routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
+        # 将权重的数据类型转换回模型的主数据类型(如bfloat16)以提高效率
+        routing_weights = routing_weights.to(hidden_states.dtype)
+
+        # --- 步骤 2: 计算并存储负载均衡损失 ---
+        # 这是MoE训练的灵魂：确保路由器不会“偏心”，公平地使用所有专家
+        router_probs_full = F.softmax(router_logits, dim=-1, dtype=torch.float)
+        avg_expert_prob = router_probs_full.mean(dim=0) # 每个专家在所有Token上的平均被选择概率
+        expert_mask_for_lb = F.one_hot(selected_experts, num_classes=self.num_experts).sum(dim=1) # 统计每个Token被分配给了哪些专家
+        avg_expert_fraction = expert_mask_for_lb.float().mean(dim=0) # 每个专家平均处理的Token比例
+        # 计算损失并乘以专家数量作为惩罚项，然后存储起来
+        self.most_recent_lb_loss = self.num_experts * torch.sum(avg_expert_prob * avg_expert_fraction)
+
+        # --- 步骤 3: 专家计算与结果融合 (矢量化高效版) ---
+        # 创建一个空的张量用于存放最终结果
+        final_hidden_states = torch.zeros_like(flat_hidden_states)
+        
+        # 这个循环只遍历Top-K次（比如2次），非常快
+        for k in range(self.top_k):
+            # 获取所有Token在第k个选择上的专家索引和权重
+            expert_indices_k = selected_experts[:, k]
+            routing_weights_k = routing_weights[:, k]
+            
+            # 这个循环遍历所有专家，但内部计算是批处理的，也非常快
+            for i in range(self.num_experts):
+                # 创建一个掩码，找到所有选择了当前专家i的Token
+                mask = expert_indices_k == i
+                if mask.any(): # 如果有任何Token选择了这个专家
+                    # 将这些Token的输入作为一个批次，交给专家i处理
+                    expert_output = self.experts[i](flat_hidden_states[mask])
+                    # 将专家的输出乘以对应的路由权重，并“添加”回最终结果张量的正确位置
+                    final_hidden_states.index_add_(0, torch.where(mask)[0], expert_output * routing_weights_k[mask].unsqueeze(1))
+        
+        # 将结果恢复成原始的(batch, sequence, hidden_size)形状并返回
+        return final_hidden_states.view(*original_shape)
+
+# --- 2. “创世纪”函数：组装、移植、改造模型 ---
+def create_moe_model():
+    """
+    这是“架构手术”函数。它负责：
+    1. 搭建一个带有MoE模块的空壳模型。
+    2. 将你预训练好的8个专家的权重精准地“移植”进去。
+    3. 冻结所有专家，只留下路由器是可训练的。
+    """
+    print("--- Starting Architectural Surgery ---")
+    
+    # 从标准模型加载配置，这是我们新模型的“基因蓝图”
+    config = AutoConfig.from_pretrained(MODEL_NAME)
+    
+    print("Step 1: Loading base model skeleton...")
+    # 加载8个专家中的任意一个，作为我们MoE模型的“骨架”
+    # 我们将使用它的词嵌入、注意力模块等非FFN部分
+    base_model = AutoModelForCausalLM.from_pretrained(
+        os.path.join(BASE_EXPERT_PATH, EXPERT_DIRS[0]),
+        torch_dtype=torch.bfloat16,
+        device_map=device
+    )
+    
+    print("Step 2: Pre-loading all expert weights into CPU memory for efficiency...")
+    # 为了提高效率，我们一次性把所有专家的权重从硬盘加载到CPU内存
+    # 我们使用 `safetensors.torch.load_file`，这是加载.safetensors文件的正确方式
+    all_experts_state_dicts = [
+        load_file(os.path.join(BASE_EXPERT_PATH, expert_dir, 'model.safetensors'), device='cpu')
+        for expert_dir in EXPERT_DIRS
+    ]
+    print("All expert weights pre-loaded.")
+
+    print("Step 3: Replacing FFNs with MoE modules and transplanting expert weights...")
+    # 遍历模型的30层
+    for layer_idx, layer in enumerate(base_model.model.layers):
+        # 在每一层，都用我们自己设计的MoEModule替换掉原来标准的LlamaMLP
+        layer.mlp = MoEModule(config).to(device, dtype=torch.bfloat16)
+        
+        # 开始“器官移植”
+        for expert_idx in range(NUM_EXPERTS):
+            # 从内存中获取当前专家的权重字典
+            expert_state_dict = all_experts_state_dicts[expert_idx]
+            # 筛选出只属于当前层级的FFN��分的权重
+            expert_mlp_weights = {
+                k.replace(f"model.layers.{layer_idx}.mlp.", ""): v
+                for k, v in expert_state_dict.items()
+                if f"model.layers.{layer_idx}.mlp." in k
+            }
+            # 将这些权重加载到MoE模块对应的专家“席位”上
+            layer.mlp.experts[expert_idx].load_state_dict(expert_mlp_weights)
+            
+    print("Step 4: Freezing all parameters except for the routers...")
+    # 这是关键策略：只训练“CEO”，不打扰已经很聪明的“专家”
+    for name, param in base_model.named_parameters():
+        if "router" not in name:
+            param.requires_grad = False
+            
+    print("\n--- Surgery Complete! MoE Model is assembled and ready for training. ---")
+    # 打印参数统计，验证我们的操作是否正确
+    trainable_params = sum(p.numel() for p in base_model.parameters() if p.requires_grad)
+    total_params = sum(p.numel() for p in base_model.parameters())
+    print(f"Total Parameters: {total_params / 1e6:.2f}M")
+    print(f"Trainable Parameters (Routers): {trainable_params}")
+    
+    return base_model
+
+# --- 3. 主流程：训练与保存 ---
+def main():
+    # 步骤 1: 调用“创世纪”函数，创造我们的模型
+    moe_model = create_moe_model()
+    
+    # 步骤 2: 创建优化器。它很聪明，只会包含那些 `requires_grad=True` 的参数（也就是路由器）
+    optimizer = optim.AdamW([p for p in moe_model.parameters() if p.requires_grad], lr=LEARNING_RATE)
+    
+    print("\n--- Preparing Simulated Mixed Dataset for Training ---")
+    # 注意：这里我们用的是完全随机的“模拟数据”，仅用于验证整个流程能跑通。
+    # 要想让路由器真正变聪明，必须用真实的、多样化的数据集替换这里。
+    mock_input_ids = torch.randint(0, moe_model.config.vocab_size, (BATCH_SIZE, SEQUENCE_LENGTH), device=device)
+    mock_labels = mock_input_ids.clone()
+    
+    print("--- Starting Router Training Loop (Optimized & Corrected) ---")
+    moe_model.train() # 将模型设置为训练模式
+    
+    start_time = time.time()
+    for epoch in range(EPOCHS):
+        optimizer.zero_grad() # 每个epoch开始前，清空上一轮的梯度
+        
+        # --- 优雅且正确的前向传播 ---
+        # 直接调用模型，Hugging Face会自动为我们处理所有复杂的内部细节（如attention mask）
+        # 同时，因为我们提供了 `labels`，它会自动计算主线任务的交叉熵损失
+        outputs = moe_model(input_ids=mock_input_ids, labels=mock_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 # 从我们设置的占位符中取出损失
+            
+        # --- 计算最终的“复合KPI”（总损失）---
+        total_loss = main_loss + LB_LOSS_COEFFICIENT * total_lb_loss
+        
+        # --- 反向传播与优化 ---
+        total_loss.backward() # 计算梯度
+        optimizer.step()      # 更新路由器权重
+        
+        # --- 打印训练日志 ---
+        if (epoch + 1) % 10 == 0:
+            elapsed_time = time.time() - start_time
+            print(f"Epoch [{epoch+1:03d}/{EPOCHS}] | 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: {elapsed_time:.2f}s")
+            start_time = time.time()
+            
+    print("\n--- Router Training Complete! ---")
+    
+    # --- 步骤 5: 将我们伟大的作品“固化”到硬盘上 ---
+    print("\n--- Phase 5: Saving the fully trained MoE model to disk ---")
+    OUTPUT_MODEL_DIR = "./SmolMoE-8x135M-Instruct-v1-Trained"
+    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...")
+    # 使用不会与Hugging Face原生配置冲突的自定义名称来保存我们的MoE参数
+    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 = AutoTokenizer.from_pretrained(MODEL_NAME) # 保存分词器
+    tokenizer.save_pretrained(OUTPUT_MODEL_DIR)
+    
+    print("\n--- Model successfully saved! ---")
+    print("You can now load this model in other scripts, but you must re-define the custom MoE classes first.")
+
+# --- 脚本入口 ---
+# 确保只有在直接运行此文件时，才会执行main函数
+if __name__ == "__main__":
+    main()