MoE

1
 class MoELayer:
2
     def __init__(self, hidden_size=768, num_experts=8, top_k=2):
3
         self.gate = Linear(hidden_size, num_experts)
4
         self.experts = ModuleList([Expert(hidden_size) for _ in
5
 range(num_experts)])
6
         self.top_k = top_k
7

8
     def forward(self, x):
9
         B, S, H = x.shape
10
         x = x.view(-1, H)
11

12
         gates = softmax(self.gate(x), dim=-1) # 权重 (B, num_experts)
13
         probs, indices = topk(gates, self.top_k, dim=-1)
14
         probs = probs / probs.sum(dim=-1, keepdim=True)
15

27 collapsed lines
16
         out = zeros_like(x)
17
         for i in range(x.shape[0]): # batch
18
             for j in range(self.top_k):
19
                 out[i] += probs[i, j] * self.experts[indices[i, j]](x[i:i+1])
20

21
         return out.view(B, S, H)
22

23
 class Expert:
24
     def __init__(self, hidden_size):
25
         self.w1 = Linear(hidden_size, hidden_size * 4)
26
         self.w2 = Linear(hidden_size * 4, hidden_size)
27
         self.act = GELU()
28

29
     def forward(self, x):
30
         return self.w2(self.act(self.w1(x)))
31

32
 class TransformerBlock:
33
     def __init__(self, hidden_size, num_heads, num_experts=8):
34
         self.attn = MultiHeadAttention(hidden_size, num_heads)
35
         self.moe = MoELayer(hidden_size, num_experts)
36
         self.norm1 = LayerNorm(hidden_size)
37
         self.norm2 = LayerNorm(hidden_size)
38

39
     def forward(self, x):
40
         x = self.norm1(x + self.attn(x))
41
         x = self.norm2(x + self.moe(x))
42
         return x