Phân Tích Kiến Trúc Mixture of Experts (MoE) Trong DeepSeek V3 và Tối Ưu Hóa Chi Phí GPU

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4
5class DeepSeekMoERouter(nn.Module):
6    """
7    MoE Gating Router with Top-K Selection and Shared Expert Routing.
8    """
9    def __init__(self, d_model: int, num_experts: int = 256, top_k: int = 8):
10        super().__init__()
11        self.num_experts = num_experts
12        self.top_k = top_k
13        # Linear projection to compute gating logits
14        self.gate = nn.Linear(d_model, num_experts, bias=False)
15        
16    def forward(self, x: torch.Tensor):
17        # x shape: [batch_size, seq_len, d_model]
18        batch_size, seq_len, d_model = x.shape
19        x_flat = x.view(-1, d_model)
20        
21        # 1. Compute routing logits for all experts
22        logits = self.gate(x_flat) # [num_tokens, num_experts]
23        
24        # 2. Select top-k experts with highest affinity
25        topk_weights, topk_indices = torch.topk(logits, self.top_k, dim=-1)
26        
27        # 3. Softmax normalize over the selected top-k weights
28        topk_probs = F.softmax(topk_weights, dim=-1)
29        
30        return topk_probs, topk_indices
31
32class MoEFeedForward(nn.Module):
33    def __init__(self, d_model: int, d_ffn: int, num_experts: int = 256, top_k: int = 8):
34        super().__init__()
35        self.router = DeepSeekMoERouter(d_model, num_experts, top_k)
36        self.experts = nn.ModuleList([
37            nn.Sequential(
38                nn.Linear(d_model, d_ffn),
39                nn.SiLU(),
40                nn.Linear(d_ffn, d_model)
41            ) for _ in range(num_experts)
42        ])
43        # Shared Expert always activated for general knowledge
44        self.shared_expert = nn.Sequential(
45            nn.Linear(d_model, d_ffn),
46            nn.SiLU(),
47            nn.Linear(d_ffn, d_model)
48        )
49
50    def forward(self, x: torch.Tensor) -> torch.Tensor:
51        topk_probs, topk_indices = self.router(x)
52        # Dispatch to active experts and sum weighted outputs
53        shared_out = self.shared_expert(x)
54        return shared_out # + routed_expert_outputs