1
import torch
2
import torch.nn as nn
3
import torch.nn.functional as F
4

5
class MixtralMoE(nn.Module):
6
    """Mixtral 的稀疏 MoE 实现"""
7

8
    def __init__(self, hidden_dim, intermediate_dim, num_experts=8, top_k=2):
9
        super().__init__()
10
        self.num_experts = num_experts
11
        self.top_k = top_k
12

13
        # 路由器：决定 token 分配给哪个专家
14
        self.router = nn.Linear(hidden_dim, num_experts, bias=False)
15

16
        # 专家网络：每个专家是一个 FFN
17
        self.experts = nn.ModuleList([
18
            FeedForward(hidden_dim, intermediate_dim)
19
            for _ in range(num_experts)
20
        ])
21

22
    def forward(self, x):
23
        """
24
        Args:
25
            x: [batch_size, seq_len, hidden_dim]
26
        """
27
        batch_size, seq_len, hidden_dim = x.shape
28

29
        # 1. 路由计算
30
        router_logits = self.router(x)  # [batch, seq, num_experts]
31

32
        # 2. 选择 Top-K 专家
33
        router_probs = F.softmax(router_logits, dim=-1)
34
        top_k_probs, top_k_indices = torch.topk(router_probs, self.top_k, dim=-1)
35

36
        # 3. 归一化选中的概率
37
        top_k_probs = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)
38

39
        # 4. 计算每个专家的输出并加权求和
40
        output = torch.zeros_like(x)
41

42
        for i in range(self.top_k):
43
            expert_idx = top_k_indices[:, :, i]  # [batch, seq]
44
            expert_weight = top_k_probs[:, :, i:i+1]  # [batch, seq, 1]
45

46
            # 对每个位置，用选中的专家计算
47
            for e in range(self.num_experts):
48
                mask = (expert_idx == e)  # [batch, seq]
49
                if mask.any():
50
                    # 获取需要用这个专家的 token
51
                    expert_input = x[mask]  # [num_tokens, hidden]
52
                    expert_output = self.experts[e](expert_input)
53

54
                    # 加权累加
55
                    output[mask] += expert_weight[mask] * expert_output
56

57
        return output
58

59

60
class FeedForward(nn.Module):
61
    """单个专家的 FFN"""
62
    def __init__(self, hidden_dim, intermediate_dim):
63
        super().__init__()
64
        self.w1 = nn.Linear(hidden_dim, intermediate_dim, bias=False)
65
        self.w2 = nn.Linear(intermediate_dim, hidden_dim, bias=False)
66
        self.w3 = nn.Linear(hidden_dim, intermediate_dim, bias=False)
67

68
    def forward(self, x):
69
        # SwiGLU 激活
70
        return self.w2(F.silu(self.w1(x)) * self.w3(x))

1
def sliding_window_attention(Q, K, V, window_size):
2
    """
3
    滑动窗口注意力
4

5
    Args:
6
        Q, K, V: [batch, heads, seq_len, head_dim]
7
        window_size: 窗口大小
8
    """
9
    batch, heads, seq_len, head_dim = Q.shape
10

11
    # 计算注意力分数
12
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(head_dim)
13

14
    # 创建滑动窗口掩码
15
    mask = torch.ones(seq_len, seq_len, dtype=torch.bool)
16
    for i in range(seq_len):
17
        # 每个 token 只能看到窗口内的 token
18
        start = max(0, i - window_size + 1)
19
        mask[i, :start] = False  # 窗口外的位置设为 False
20

21
    # 应用掩码
22
    scores = scores.masked_fill(~mask, float('-inf'))
23

24
    # Softmax 和加权求和
25
    attention = F.softmax(scores, dim=-1)
26
    output = torch.matmul(attention, V)
27

28
    return output

1
| 维度          | 标准注意力       | 滑动窗口注意力            |
2
|---------------|------------------|---------------------------|
3
| 计算复杂度    | O(n²)           | O(n × w)                  |
4
| 内存占用      | O(n²)           | O(n × w)                  |
5
| 长序列支持    | 差              | 好                        |
6
| 远距离依赖    | 直接建模        | 通过多层间接建模          |
7
| 实现复杂度    | 简单            | 中等                      |
8

9
Mixtral 的配置：
10
• 窗口大小 w = 4096
11
• 支持序列长度 n = 32768
12
• 计算量减少约 8 倍

1
| 模型          | 参数/激活   | MMLU        | HumanEval       |
2
|---------------|-------------|-------------|------------------|
3
| LLaMA 2 7B    | 7B / 7B     | 45.3%       | 12.8%           |
4
| Mistral 7B    | 7B / 7B     | 62.5%       | 26.2%           |
5
| LLaMA 2 13B   | 13B / 13B   | 54.8%       | 18.9%           |
6
| LLaMA 2 34B   | 34B / 34B   | 62.6%       | 27.8%           |
7
| LLaMA 2 70B   | 70B / 70B   | 69.8%       | 32.3%           |
8
| Mixtral 8x7B  | 47B / 13B   | 70.6%       | 40.2%           |
9
| GPT-3.5       | -           | 70.0%       | 48.1%           |
10

11
关键发现：
12
• Mixtral 在 MMLU 上超越 LLaMA 2 70B
13
• 在代码任务（HumanEval）上表现优异
14
• 推理成本仅为 LLaMA 2 70B 的约 1/5

1
| 模型          | 总参数      | 激活参数    | 相对推理成本    |
2
|---------------|-------------|-------------|------------------|
3
| LLaMA 2 70B   | 70B         | 70B         | 5.4x            |
4
| Mixtral 8x7B  | 47B         | 13B         | 1x（基准）      |
5
| LLaMA 2 13B   | 13B         | 13B         | 1x              |
6
| Mistral 7B    | 7B          | 7B          | 0.54x           |
7

8
MoE 的效率优势：
9
• 47B 总参数，但只激活 13B
10
• 推理成本接近 13B 稠密模型
11
• 性能却接近 70B 稠密模型
12
• 性价比极高

1
| 语言          | LLaMA 2 70B | Mixtral     | 提升            |
2
|---------------|-------------|-------------|------------------|
3
| 英语          | 69.8%       | 70.6%       | +0.8%           |
4
| 法语          | 63.4%       | 69.2%       | +5.8%           |
5
| 德语          | 62.7%       | 68.4%       | +5.7%           |
6
| 西班牙语      | 64.1%       | 69.8%       | +5.7%           |
7
| 意大利语      | 61.8%       | 68.1%       | +6.3%           |
8

9
Mixtral 在非英语语言上表现更优，可能因为：
10
• 更大的模型容量
11
• 专家可以专门处理不同语言
12
• 训练数据包含更多多语言内容

1
| 维度          | 稠密模型         | MoE 模型                  |
2
|---------------|------------------|---------------------------|
3
| 参数效率      | 线性             | 更高（容量大）            |
4
| 计算效率      | 随参数增长       | 随激活参数增长            |
5
| 内存占用      | 随参数增长       | 需要存储所有专家          |
6
| 训练稳定性    | 稳定             | 需要负载均衡              |
7
| 推理延迟      | 随参数增长       | 接近激活参数规模          |
8
| 硬件友好度    | 高               | 中（需要优化）            |