[AI学习笔记]DeepSeek混合专家系统（MoE）架构深度解析

# 传统密集模型的问题演示
class DenseModel(nn.Module):
    def __init__(self, dim=1024):
        self.layer = nn.Sequential(
            nn.Linear(dim, 4*dim),
            nn.GELU(),
            nn.Linear(4*dim, dim)
        )  # 参数量: 1024*(4096+1024)=5.2M/layer

# DeepSeek MoE对比
class MoEModel(nn.Module):
    def __init__(self, num_experts=8):
        self.experts = nn.ModuleList([Expert(256) for _ in range(num_experts)])
        self.gate = nn.Linear(1024, num_experts)  # 参数量: 1024*8 + 8*256*2=0.1M/layer

class DynamicDualGate(nn.Module):
    def __init__(self, dim, num_experts):
        self.main_gate = nn.Linear(dim, num_experts)
        self.entropy_bias = nn.Parameter(torch.zeros(num_experts))
        
    def forward(self, x):
        logits = self.main_gate(x)
        entropy = -torch.sum(F.softmax(logits, dim=-1) * F.log_softmax(logits, dim=-1)
        logits += self.entropy_bias * entropy.unsqueeze(-1)
        return F.softmax(logits, dim=-1)

# 基于PyTorch Distributed的专家并行
def expert_forward(inputs, expert_idx):
    world_size = dist.get_world_size()
    # 按专家索引分配设备
    device_idx = expert_idx % world_size
    inputs = inputs.to(f'cuda:{device_idx}')
    # 本地专家计算
    output = experts[expert_idx](inputs)
    # 全局结果收集
    output_list = [torch.empty_like(output) for _ in range(world_size)]
    dist.all_gather(output_list, output)
    return torch.cat(output_list, dim=0)

# 调用示例
outputs = []
for idx in active_experts:
    out = expert_forward(inputs, idx)
    outputs.append(out)
final_output = sum(outputs) / len(outputs)

class LoadBalancer:
    def __init__(self, num_experts):
        self.usage = torch.zeros(num_experts)
        self.alpha = 0.9  # 平滑系数
        
    def update(self, expert_indices):
        batch_usage = torch.bincount(expert_indices, minlength=len(self.usage))
        self.usage = self.alpha * self.usage + (1 - self.alpha) * batch_usage
        
    def penalty_loss(self):
        mean_usage = self.usage.mean()
        var_loss = torch.var(self.usage) 
        skew_loss = torch.max(self.usage) / torch.min(self.usage)
        return 0.1 * var_loss + 0.05 * skew_loss

# 训练循环中
optimizer.zero_grad()
loss = task_loss + balancer.penalty_loss()
loss.backward()
optimizer.step()

#!/bin/bash
# DeepSeek MoE集群启动脚本
NUM_NODES=128
GPUS_PER_NODE=8

torchrun \
    --nnodes=$NUM_NODES \
    --nproc_per_node=$GPUS_PER_NODE \
    --rdzv_backend=c10d \
    --rdzv_endpoint=coordinator:29500 \
    train_moe.py \
    --batch_size 4096 \
    --expert_parallel_degree=32 \
    --num_experts=256 \
    --gate_type=dynamic \
    --capacity_factor=1.2

class Dashboard:
    metrics = {
        'expert_usage': [], 
        'throughput': [],
        'loss': []
    }
    
    def update(self, batch_data):
        self.metrics['expert_usage'].append(balancer.usage.numpy())
        self.metrics['throughput'].append(batch_data['tokens'] / batch_data['time'])
        self.metrics['loss'].append(loss.item())
        
    def visualize(self):
        plt.figure(figsize=(12,4))
        plt.subplot(131)
        plt.plot(self.metrics['expert_usage'])
        plt.title('Expert Utilization')
        # 其他子图类似...

class MultimodalMoE(nn.Module):
    def forward(self, text, image):
        # 文本分支
        text_feat = self.text_moe(text)
        # 图像分支
        image_feat = self.vision_moe(image)
        # 跨模态融合
        fused = self.cross_gate(torch.cat([text_feat, image_feat]))
        return self.head(fused)

# 输入样例
input_text = "图示化学结构式的名称是？"
input_image = load_image("benzene.png")
output = model(input_text, input_image)  # 输出：苯环结构