train/train_moe_router.py

# ==============================================================================
#                      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()