本文基于博客 LLM MOE 的进化之路，介绍 InfiniLM 中的 MOE 实现。
第一个环节是顺着这个博客的逻辑简单介绍三种 MOE 实现。
为什么要有 MOE 呢？随着模型的规模增大，泛化能力增强，现在的模型普遍变得越来越不可解释了。CNN 的时代还有一些研究讨论神经网络的每个层分别在产生最终结果的过程中发挥了什么作用，对于 Transformer 模型，则很少看到这样的讨论了。但是，人们还是可以通过观察矩阵中数值的大小来直观感受到矩阵的某个部分对于最终结果的贡献大小。通过这种观察，大家发现了模型激活的稀疏性。所谓稀疏，其实是不均匀的意思。也就是说，往往一次推理并不是所有的权重都对结果同样重要。
但是，由于推理的计算方式，无论重要的还是不重要的权重都要付出同样多的计算开销，这显然是不够经济的。所以产生了 MOE，也就是混合专家结构。MOE 的思路很简单，既然不是所有参数都会在计算中用到，那么就在训练时直接把权重划分到几个并列的结构，然后每次只用其中的一些来计算。关于稀疏现象和稀疏模型的有效训练，可以看本次训练营大模型前沿技术中的介绍。今天我们只从结果上介绍 MOE 模型的结构和实现。
一个典型的 MOE 结构可以看博客上的图：
实际上就是把一个正常的结构缩小，同时复制 n 份，并添加一个门控网络。门控网络的功能是根据输入为每个专家赋权，每个专家都对输入做各自的处理，最后做个加权平均作为总体的结果。
显然单纯做这个加权平均完全没体现稀疏性。但是这个思路是很通顺的，既然一开始发现一个大矩阵里面不均匀，于是把它分开。现在有了门控之后，专家之间的不均匀性就是可量化的了。只需要把最重要的哪些专家挑出来，不重要的专家直接不激活。如图所示：
这个图说的是所有类似 llama 的模型，MOE 都坐在 FFN 层。模型的参数里会指定 FFN 层被复制了多少份，以及每次实际计算应该激活多少份：
block_count······················ 32
context_length··················· 32768
embedding_length················· 4096
feed_forward_length·············· 14336
attention.head_count············· 32
attention.head_count_kv·········· 8
rope.freq_base··················· 1e6
attention.layer_norm_rms_epsilon· 1e-5
expert_count····················· 8
expert_used_count················ 2
vocab_size······················· 32000
rope.dimension_count············· 128
这里的 expert_count 和 expert_used_count 说的就是这个模型有每层有 8 份完整的 FFN 参数，每次应该激活其中 2 个。
最后是这个带共享专家的稀疏 MOE，实际上就是有几个不受门控管制的高级专家，所有 token 都用，其他专家 m 选 n，最后再一起加权。由于实现上并没有特别的难度，我这里也没有能跑的模型可以演示，就不讲了。
直接来看代码：
#[derive(Clone, Debug)]
pub struct LlamaMeta {
pub nexp: usize,
pub nexp_use: usize,
..
}元信息里需要保存这两个参数。
# [derive(Clone)]
pub struct BlkStorage<T> {
pub attn_norm: T,
pub attn_qkv: T,
pub attn_qkv_bias: T,
pub attn_o: T,
pub ffn_norm: T,
pub ffn_gate_inp: T,
pub ffn_gate_up: T,
pub ffn_down: T,
}层的存储里，可以看到多了 1 项 gate_inp，这就是那个门控网络。另外现在 gate_up 和 down 在 MOE 的情况下是带专家的，比稠密多一个维度。在 Mixtral 里是这样的：
blk.x.ffn_gate_inp.weight·······F16 [4096, 8]
blk.x.ffn_gate_up_exps.weight···F16 [4096, 28672, 8]
blk.x.ffn_down_exps.weight······F16 [14336, 4096, 8]
门控网络可以输入一个 hidden size 长的矩阵或者说一组向量，计算出每个专家的权重，所以它形状是专家数乘隐向量长度。接下来通过一个 softmax 和 topk 的融合算子，计算出要使用的专家和权重。这种融合比较像你们应该已经写过的 random sample。由于这个权重会被传递给 down 计算的环节，这里必然涉及一次 device to host 拷贝。其实硬说的话这个拷贝也是可规避的，但是我们写的比较简单，直接在门控网络之后就把结果拷贝到 host，用 cpu 做的 topk 融合计算。完整代码见 llama/compute.rs。