【LLM架构】动手实现MoE架构

前文

前面我们对MoE架构的原理进行了解释，这里简单回顾一下：

MOE主要由两个关键点组成：

1. 将传统Transformer中的FFN（前馈网络层）替换为多个稀疏的专家层（Sparse MoE layers）。每个专家本身是一个独立的神经网络，实际应用中，这些专家通常是前馈网络 (FFN)，但也可以是更复杂的网络结构。
2. 门控网络或路由：此部分用来决定输入的token分发给哪一个专家。

这里我们会从头到尾实现一层MoE架构，里面会涉及较多tensor的操作，我们也会一起讲解，所以内容可能比较多，请见谅！

FFN（前馈网络层）

可能有对FFN（前馈网络层）不太熟悉的小伙伴可以看一下下面的代码及图例，很简单就是一个我们平时常见的结构。

专家模型

专家模型其实就是一个多层感知机MLP。

nn.Sequential详解

在给定的代码中，nn.Sequential 的作用是将多个神经网络层按顺序组合成一个子网络，简化模型的构建和前向传播过程。具体来说：

1. 层顺序组合：nn.Sequential 会按照添加顺序依次堆叠层（如 nn.Linear、nn.ReLU 等），确保前一层输出直接作为后一层输入。
2. 自动前向传播：在 forward 方法中直接调用 self.net(x) 即可完成整个子网络的前向计算，无需手动逐层传递数据。
3. 代码简洁性：通过 nn.Sequential 包装层，避免了显式定义 forward 方法中逐层操作的繁琐，尤其适合结构简单的网络。

例如，代码中的 self.net 实际上等价于一个包含两个线性层、ReLU 激活和 Dropout 的顺序连接网络，其前向传播会依次执行这些操作。

在PyTorch中，nn.Sequential容器可以添加大多数继承自nn.Module的子模块（如线性层、卷积层、激活函数等），但存在以下限制：

1. 可以添加的子模块类型

• 标准层：如nn.Linear、nn.Conv2d、nn.ReLU、nn.MaxPool2d等。
• 自定义层：需继承nn.Module并实现forward方法。
• 嵌套容器：如另一个nn.Sequential或nn.ModuleList。
• 命名层：可通过add_module或OrderedDict为层指定名称。

2. 不能添加或受限的情况

• 动态结构：如条件分支（if-else）、循环结构或根据输入动态选择路径的逻辑，需自定义forward方法。
• 跳跃连接（Residual Connections）：需手动实现前向传播中的残差相加。
• 非线性交互：如层间非顺序的复杂组合（如x = layer1(x) + layer2(x)），需自定义forward。
• 无forward方法的模块：若子模块未正确实现forward，会导致运行时错误。

3. 示例说明

• 允许：线性堆叠层

点击链接【LLM架构】动手实现MoE架构阅读全文