Mixtral 8x7B 详解：稀疏专家混合模型的架构与源码分析

解读 Mistral AI 发布的 Mixtral 8x7B 稀疏专家混合模型：支持 32K 上下文、多语言与代码生成能力突出，并顺着 transformers 源码逐层分析 MixtralConfig、MixtralModel 的初始化与前向计算流程。

一、模型结构分析

前言

2023年12月11日，Mistral AI团队发布了一款高质量的稀疏专家混合模型Mixtral 8x7B。

Mistral AI继续致力于向开发者社区提供最优秀的开放模型。在人工智能领域向前发展，需要采取超越重用众所周知的架构和训练范式的新技术路径。最重要的是，它需要让社区从原创模型中受益，以促进新的发明和用途。

Mixtral具有以下特点：

• 优雅地处理32k标记的上下文。
• 支持英语、法语、意大利语、德语和西班牙语。
• 在代码生成方面表现出色。
• 可以微调为一个遵循指令的模型，在MT-Bench上达到8.3的分数。

在 transformers 仓库中可以看到 mixtral 的源码，首先是 MixtralModel 类，继承自 PreTrainedModel ，这个类是所有模型的基类，包含了一些通用的方法，比如保存模型、加载模型、初始化权重等。具体目录是：src\transformers\models\mixtral\modeling_mixtral.py

继承关系为：MixtralModel -> MixtralPreTrainedModel -> PreTrainedModel

MixtralConfig

MixtralConfig 类继承自 PretrainedConfig ，这个类是所有配置类的基类，包含了一些通用的方法，比如保存配置、加载配置、初始化配置等。具体路径在 transformers 仓库的 src\transformers\models\mixtral\configuration_mixtral.py目录下。

可以使用如下代码直接创建模型的config对象：

config = MixtralConfig()

MixtralModel

MixtralModel 初始化

如果你看过我上一篇 LLaMA开源大模型源码分析！的话，就会发现这里的初始化和llama模型的初始化非常相似，都是先初始化embed_tokens，然后初始化layers，最后初始化norm。

• 设置了模型的两个属性:padding_idx（用于指定填充标记的索引），vocab_size（词汇表的大小）
• 初始化了模型的嵌入层、解码器层、归一化层
• 嵌入层（nn.Embedding）：模型使用嵌入层将输入的标记映射成密集的向量表示。
• 解码器层（nn.ModuleList()）：模型包含多个解码器层，这些层都是由 MixtralDecoderLayer 定义
• 归一化层 MixtralRMSNorm：归一化层使用的是 Root Mean Square Layer Normalization（RMS Layer Norm），和llama使用的是一样的。
• 设置了是否使用 gradient_checkpoint 主要是用来节省显存
• 调用 post_init() 完成一些初始化和准备检查的代码

  class MixtralModel(MixtralPreTrainedModel):
      """
      Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`MixtralDecoderLayer`]
  
      Args:
          config: MixtralConfig
      """
  
      def __init__(self, config: MixtralConfig):
          super().__init__(config)
          self.padding_idx = config.pad_token_id
          self.vocab_size = config.vocab_size
  
          self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
          self.layers = nn.ModuleList(
              [MixtralDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
          )
          self._attn_implementation = config._attn_implementation
          self.norm = MixtralRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
  
          self.gradient_checkpointing = False
  # Initialize weights and apply final processing
          self.post_init()

可以看一下 post_init() 的代码，主要是初始化权重和gradient_checkpointing相关的一些事情。该方法在PreTrainedModel基类中，transformers中所有模型基本都继承这个类。

def post_init(self):
    """
    A method executed at the end of each Transformer model initialization, to execute code that needs the model's
    modules properly initialized (such as weight initialization).
    """
    self.init_weights()
    self._backward_compatibility_gradient_checkpointing()

MixtralModel Forward

forward 部分的代码有点长，但其实大部分都是张量并行或者是节省显存相关的代码，对于理解模型结构来说可以直接忽略。

首先进来就是把 inputs_ids 进行向量化，然后拿到 hidden_states 。然后是存起来所有的hidden_states 进入 decoder_layer 再拿一个 hidden_states，作为下一轮 decoder_layer 的 hidden_states 输入，最后给 hidden_states norm一下。如下代码所示：

# 向量化
inputs_embeds = self.embed_tokens(input_ids)
hidden_states = inputs_embeds

for decoder_layer in self.layers:
    #存起来所有的 hidden_states
    if output_hidden_states:
        all_hidden_states += (hidden_states,)
    # 这里是decoder_layer 的forward
    layer_outputs = decoder_layer(
        hidden_states,
        attention_mask=attention_mask,
        position_ids=position_ids,
        past_key_value=past_key_values,
        output_attentions=output_attentions,
        output_router_logits=output_router_logits,
        use_cache=use_cache,
    )
    # # 再拿一个 hidden_states，作为下一轮 decoder_layer 的 hidden_states 输入
    hidden_states = layer_outputs[0]

# norm 一下
hidden_states = self.norm(hidden_states)

MixtralDecoderLayer

MixtralDecoderLayer 初始化

好，来到了 moe 模型和 llama 模型最大区别的地方了，Mixtral 使用 MixtralSparseMoeBlock 模块代替了原有的 MLP 层， MLP 层还是在的，待会在后面我们再说。先来看初始化部分 DecoderLayer 做了什么事情。

• hidden_size : 也就是在上面说的输入输出。

• self_attn : 别看它写这么多啊，其实就是选一下用什么 attention 。看见大写字母不要怕，直接点进去看看怎么个事！

  MIXTRAL_ATTENTION_CLASSES = {
      "eager": MixtralAttention,
      "flash_attention_2": MixtralFlashAttention2,
      "sdpa": MixtralSdpaAttention,
  }

• block_sparse_moe : moe稀疏矩阵，这个待会后面再说，输入输出都是 hidden_size 大小。

• input_layernorm : MixtralRMSNorm 层，输入时候的norm

• post_attention_layernorm : 丢入稀疏矩阵 block_sparse_moe 之前的操作。

  class MixtralDecoderLayer(nn.Module):
      def __init__(self, config: MixtralConfig, layer_idx: int):
          super().__init__()
          self.hidden_size = config.hidden_size# 隐藏层的大小
  
          self.self_attn = MIXTRAL_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx)# 自注意力机制
  
          self.block_sparse_moe = MixtralSparseMoeBlock(config)# 稀疏混合块
          self.input_layernorm = MixtralRMSNorm(config.hidden_size, eps=config.rms_norm_eps)# 输入层归一化
          self.post_attention_layernorm = MixtralRMSNorm(config.hidden_size, eps=config.rms_norm_eps)# 注意力之后的层归一化

MixtralDecoderLayer Forward

首先复制一份 hidden_states 给 residual。然后 hidden_states 进入 input_layernorm 进行norm。

然后进入 self_attn 进行 attention 操作，拿到 hidden_states、self_attn_weights、present_key_value。

而后 hidden_states 和 residual 相加，得到 hidden_states。此时再复制一份 residual 。然后 hidden_states 进入 post_attention_layernorm 进行norm。

来了，来了！这里 hidden_states 进入稀疏矩阵 block_sparse_moe 得到 hidden_states, router_logits ，hidden_states 和 residual 相加，得到 hidden_states。最后输出 hidden_states。

residual = hidden_states
hidden_states = self.input_layernorm(hidden_states)

hidden_states, self_attn_weights, present_key_value = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_value=past_key_value,
            output_attentions=output_attentions,
            use_cache=use_cache,
        )

hidden_states = residual + hidden_states

residual = hidden_states
hidden_states = self.post_attention_layernorm(hidden_states)
hidden_states, router_logits = self.block_sparse_moe(hidden_states)
hidden_states = residual + hidden_states

outputs = (hidden_states,)

if output_attentions:
    outputs += (self_attn_weights,)

if use_cache:
    outputs += (present_key_value,)

if output_router_logits:
    outputs += (router_logits,)

return outputs

MixtralAttention

我们先来看 Attention 部分嗷，稀疏矩阵留到最后压轴再看。

MixtralAttention 初始化

好好好，首先映入眼帘的还是 Attention Is All You Need ，不忘初心，可以可以！

先来看 init 部分叭。

• layer_idx : 这个就是第几个 DecoderLayers 层。不用关心。
• attention_dropout : 用于dropout的概率。
• hidden_size : 输入输出大小。
• num_attention_heads : 多头注意力的头数。
• head_dim : 多头注意力的维度 self.hidden_size // self.num_heads，和transformers中的一样。
• num_key_value_heads : 用于key和value的头数。
  其他的参数都在 MixtralConfig 中有默认值，可以直接使用，也可以直接去MixtralConfig的源码中看具体的解释，这里就不再多说。

再往下就是 q_proj、 k_proj 、v_proj、 o_proj 四个矩阵（全连接层），耳熟能详了。

class MixtralAttention(nn.Module):
    """
    Multi-headed attention from 'Attention Is All You Need' paper. Modified to use sliding window attention: Longformer
    and "Generating Long Sequences with Sparse Transformers".
    """

    def __init__(self, config: MixtralConfig, layer_idx: Optional[int] = None):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        if layer_idx is None:
            logger.warning_once(
                f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "
                "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
                "when creating this class."
            )

        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.head_dim = self.hidden_size // self.num_heads
        self.num_key_value_heads = config.num_key_value_heads
        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
        self.max_position_embeddings = config.max_position_embeddings
        self.rope_theta = config.rope_theta
        self.is_causal = True
        self.attention_dropout = config.attention_dropout

        if (self.head_dim * self.num_heads) != self.hidden_size:
            raise ValueError(
                f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
                f" and `num_heads`: {self.num_heads})."
            )
        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)

        self.rotary_emb = MixtralRotaryEmbedding(
            self.head_dim,
            max_position_embeddings=self.max_position_embeddings,
            base=self.rope_theta,
        )

MixtralAttention Forward

这里的 forward 函数就是 Attention 的核心部分了，我们来一点一点看。

注意：其中有关于张量并行或者显存节省的部分我就直接省略了，直接看主要代码。这个笔记主要是分析mixtral的模型结构，并不讨论如何节省显存。

首先获取 batch_size 和 seq_len ，然后把 hidden_states 丢入 q_proj、 k_proj 、v_proj 三个矩阵，得到 query_states、 key_states 、value_states 。然后把 query_states、 key_states 、value_states reshape 为下一步计算做准备。

获取 kv_seq_len ，其实我觉得这步挺多余的，因为 kv_seq_len 就等于 self.num_key_value_heads 。

将旋转位置嵌入应用于查询和键张量。使用了旋转位置嵌入的余弦和正弦部分，将它们与查询和键张量相乘，并将结果相加，从而实现旋转位置嵌入的效果。

key_states和value_states重复self.num_key_value_groups次。然后，使用torch.matmul()函数计算query_states和转置后的key_states之间的矩阵乘法。最后，将结果除以math.sqrt(self.head_dim)进行归一化。

然后softmax 和 dropout。然后 attn_weights 和 value_states 相乘，把 attn_output reshape 为下一步计算做准备，最后把 attn_output 丢入 o_proj ，然后return就行了。

# 获取 batch_size 和 seq_len
bsz, q_len, _ = hidden_states.size()

# 把 hidden_states 丢入 q_proj、k_proj、v_proj
query_states = self.q_proj(hidden_states)
key_states = self.k_proj(hidden_states)
value_states = self.v_proj(hidden_states)

# 把 q_proj、k_proj、v_proj 的输出 reshape 为下一步计算做准备
query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)

# 获取 kv_seq_len，其实我觉得这步挺多余的，因为 kv_seq_len 就等于 self.num_key_value_heads
kv_seq_len = key_states.shape[-2]

# 将旋转位置嵌入应用于查询和键张量。使用了旋转位置嵌入的余弦和正弦部分，将它们与查询和键张量相乘，并将结果相加，从而实现旋转位置嵌入的效果
cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)

# 首先，它将key_states和value_states重复self.num_key_value_groups次。然后，使用torch.matmul()函数计算query_states和转置后的key_states之间的矩阵乘法。最后，将结果除以math.sqrt(self.head_dim)进行归一化
key_states = repeat_kv(key_states, self.num_key_value_groups)
value_states = repeat_kv(value_states, self.num_key_value_groups)
attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)

# softmax + dropout
attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)

# 然后 attn_weights 和 value_states 相乘
attn_output = torch.matmul(attn_weights, value_states)

# 然后把 attn_output reshape 为下一步计算做准备
attn_output = attn_output.transpose(1, 2).contiguous()
attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
# 最后把 attn_output 丢入 o_proj
attn_output = self.o_proj(attn_output)

# 返回 attn_output、attn_weights、past_key_value
return attn_output, attn_weights, past_key_value

MixtralSparseMoeBlock

来了，来了。MoE模型的核心，稀疏矩阵！

MixtralSparseMoeBlock 初始化

首先来看看在初始化中，init做了什么事情。

• hidden_dim : 输入输出维度大小。
• ffn_dim : MLP 层的维度大小。
• num_experts : 本地专家的数量。
• top_k : 选择的专家数量。
• gate : 门控层，输入是 hidden_dim ，输出是 num_experts 。
• experts : 专家层，八个 MixtralBLockSparseTop2MLP 模块。（就是八个原来的MLP层）

  class MixtralSparseMoeBlock(nn.Module):
  
      def __init__(self, config):
          super().__init__()
          self.hidden_dim = config.hidden_size
          self.ffn_dim = config.intermediate_size
          self.num_experts = config.num_local_experts
          self.top_k = config.num_experts_per_tok
  
  # gating
          self.gate = nn.Linear(self.hidden_dim, self.num_experts, bias=False)
  
          self.experts = nn.ModuleList([MixtralBLockSparseTop2MLP(config) for _ in range(self.num_experts)])

MixtralSparseMoeBlock Forward

• 首先，输入的隐藏状态hidden_states经过重塑，以适应后续处理。
• 使用门控层gate计算出每个隐藏状态对于各个专家的重要程度，得到router_logits。
• 对router_logits应用softmax函数，得到路由权重routing_weights。
• 从routing_weights中选出最相关的top_k个专家，并进行归一化。
• 初始化最终的隐藏状态final_hidden_states。
• 对每个专家进行遍历，根据专家掩码expert_mask选出分配给当前专家的隐藏状态，经过专家层处理后，将结果累加到最终隐藏状态中。
• 最后，将最终隐藏状态的形状重塑回原始形状，并返回。
  看完了稀疏矩阵的数据流向，现在你还觉得MoE模型在推理的之后只有两个模型在运行嘛？哈哈哈，其实就是八个MLP层作为专家模型，实际上所有的八个MLP层都是在运行的。

# 首先获取隐藏状态的维度信息
batch_size, sequence_length, hidden_dim = hidden_states.shape
# 将隐藏状态的形状重塑为二维，便于后续处理
hidden_states = hidden_states.view(-1, hidden_dim)

# router_logits用于计算每个专家对每个隐藏状态的重要程度
router_logits = self.gate(hidden_states)

# 使用softmax函数计算路由权重，这些权重决定每个隐藏状态分配给每个专家的比例
routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)
# 选择top_k个最相关的专家
routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
# 对路由权重进行归一化处理
routing_weights /= routing_weights.sum(dim=-1, keepdim=True)

# 将路由权重转换回输入数据类型
routing_weights = routing_weights.to(hidden_states.dtype)

# 初始化最终隐藏状态
final_hidden_states = torch.zeros(
    (batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device
)

# 生成专家掩码，用于确定哪些隐藏状态分配给哪些专家
expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0)

# 遍历所有的专家
for expert_idx in range(self.num_experts):
# 获取当前专家的处理层
    expert_layer = self.experts[expert_idx]
# 找出选中当前专家的隐藏状态索引
    idx, top_x = torch.where(expert_mask[expert_idx])

# 如果没有隐藏状态被分配给当前专家，则继续下一个专家
    if top_x.shape[0] == 0:
        continue

# 将索引转换为列表形式，以便高效处理
    top_x_list = top_x.tolist()
    idx_list = idx.tolist()

# 获取并处理当前专家应处理的隐藏状态
    current_state = hidden_states[None, top_x_list].reshape(-1, hidden_dim)
    current_hidden_states = expert_layer(current_state) * routing_weights[top_x_list, idx_list, None]

# 将计算结果累加回最终隐藏状态中
    final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype))

# 将最终隐藏状态的形状重塑回原始的三维形状
final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)

# 返回最终的隐藏状态和路由逻辑结果
return final_hidden_states, router_logits

MixtralBLockSparseTop2MLP

这个就是所谓的专家模型，其实就是原来的MLP层而已。

首先初始胡三个线性层和一个激活层，然后就是前向传播部分了。hidden_states 经过第一个线性层，然后经过激活层，再与经过第三个线性层的hiden_states相乘，得到current_hidden_states。

然后current_hidden_states经过第二个线性层，最后返回current_hidden_states。

class MixtralBLockSparseTop2MLP(nn.Module):
    def __init__(self, config: MixtralConfig):
        super().__init__()
        self.ffn_dim = config.intermediate_size
        self.hidden_dim = config.hidden_size

        self.w1 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False)
        self.w2 = nn.Linear(self.ffn_dim, self.hidden_dim, bias=False)
        self.w3 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False)

        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, hidden_states):
        current_hidden_states = self.act_fn(self.w1(hidden_states)) * self.w3(hidden_states)
        current_hidden_states = self.w2(current_hidden_states)
        return current_hidden_states

二、Mixtral 模型对比

https://mp.weixin.qq.com/s/lolG9Gvj_UGI4oz5oie8iw

Mistral AI 公司发布全球首款MoE(Mixture-of-Experts)大模型——Mixtral-8x7B 以来，就在AI界引起了不小的轰动，从一众科技自媒体的报道中我注意到了一个关键信息点：比Llama-2 70B具有更少的参数 ，却有更高的精度 。这一点燃起了我的兴趣，故特来学习一下Mixtral 8x7B 相对于Llama 2 70B有何不同。还是老样子

• paper ：https://arxiv.org/pdf/2401.04088.pdf
• code ：https://github.com/mistralai/mistral-src
  首先，通过Mistral AI 公司的主页我发现他一共发布了两个模型：Mistral 7B 和 Mixtral-8x7B ，后者为基于前者的MoE模型。从其公布的测试结果可以发现Mistral 7B 以7B的参数量在所有benchmarks超越了Llama-2 13B 并且与Llama-2 34B性能相当

而使用MoE策略的 Mixtral-8x7B 模型则以46.7B参数量，在多数benchmarks上超越Llama 2 70B模型。

如此优异的表现，本文就来看看这两个模型相对于Llama 2做了哪些改变，以及相对于Llama 2 这两个模型的参数量和FLOPs。

1 Mistral 7B模型

Mistral 7B模型与Llama 2 7B模型结构整体上是相似的，其结构参数如下所示

具体而言，就是存在以下几点差异：

• 对于Attention部分使用GQA (Group Query Attention)来计算注意力机制，其中Q的头数为32，而KV 的头数为8，换句话说就是每4组Q共享一组KV。这一点与Llama 2 不同，Llama 2 是在34B和70B中才使用了GQA，在7B中依然使用的是MHA(Multi-Head-Attention)
• 使用SWA(Sliding Window Attention) 。GQA和SWA叠加来降低显存占用提高推理速度。
• 增大FeedForward HiddenDim的值，由Llama-2 7B的11008 ，改为14336
  GQA和更改FFN HiddenDim的值 这两个改动都比较容易理解，那么接下来就主要来看看SWA(Sliding Window Attention)的原理和实现细节

SWA(Sliding Window Attention)

Mistral 使用了GQA和SWA两种方法来加速计算Attention，GQA在Llama 2详解的文章中说明过，这里主要讲解一下SWA。我们知道在Attention的计算一般是Q 与shape为[bst, multi-head,seq_len, head_dim] 的KV进行注意力计算，其中seq_len为已处理所有tokens总数，GQA在多头上做文章使得多组Q共享一组KV；而SWA则是在seq_len这个维度做文章，不在将Q与所有seq-len的KV “直接“计算注意力，而是只与Sliding Window Size个KV“直接“计算注意力，如下示意图，为Sliding Window Size为3的情况

举个例子，在on单词所对应的token计算Attention时，对于普通Attention $Q_{on}$ 可以与前面所有单词对应的 KV 计算Attention，而对于SWA， $Q_{on}$ 只能直接与 $Q_{on}$ 、 $Q_{sat}$ 、 $Q_{cat}$ 计算。

我们知道在LLM推理时，一般分为prompting 和 generation两个阶段，为了满足SWA，prompting阶段可以通过一个mask的掩码操作实现，如下

我们知道在LLM推理时，一般分为prompting 和 generation两个阶段，为了满足SWA，prompting阶段可以通过一个mask的掩码操作实现，如下

if input_ids.shape[1] > 1:
# seqlen推理时在prompt阶段为n，在generation阶段为1
    seqlen = input_ids.shape[1]
# mask在推理时也只在prompt阶段有,
#定义一个全1方阵
    tensor = torch.full((seqlen, seqlen),fill_value=1)
# 上三角部分全为0
    mask = torch.tril(tensor, diagonal=0).to(h.dtype)
# make the mask banded to account for sliding window
# 这里代码diagonal应该等于(-self.args.sliding_window+1)才能满足window size为
# self.args.sliding_window，这应该是官方代码的一个小bug？
    mask = torch.triu(mask, diagonal=-self.args.sliding_window)
    mask = torch.log(mask)
"""
举个例子，tensor.shape ： [10,10]
self.args.sliding_window = 5,则mask为
tensor([[1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
        [1, 1, 1, 0, 0, 0, 0, 0, 0, 0],
        [1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
        [1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
        [1, 1, 1, 1, 1, 1, 0, 0, 0, 0],
        [0, 1, 1, 1, 1, 1, 1, 0, 0, 0],
        [0, 0, 1, 1, 1, 1, 1, 1, 0, 0],
        [0, 0, 0, 1, 1, 1, 1, 1, 1, 0],
        [0, 0, 0, 0, 1, 1, 1, 1, 1, 1]])
"""

而在generation阶段，因为是自回归生成所以mask起不到作用，那此时mistral则使用了RotatingBufferCache来实现此操作，具体而言，就是采用一种循环右移的存储方式，剔除离得远的KV，保存靠近的KV 。

如上图展示了一个Window Size为4的Cache，循环右移的写Cache的示意图。

RotatingBufferCache代码实现如下

# The cache is a rotating buffer
# positions[-self.sliding_window:] 取最后w个位置的索引，取余
# [None, :, None, None]操作用于扩维度[1,w,1,1]
scatter_pos = (positions[-self.sliding_window:] % self.sliding_window)[None, :, None, None]
# repeat操作repeat维度 [bsz, w, kv_head, head_dim]
scatter_pos = scatter_pos.repeat(bsz, 1, self.n_kv_heads, self.args.head_dim)
# src取[:,-w,:,:] 所以src.shape=[bsz,w,kv_head,head_dim]
# 根据scatter_pos作为index 将src写入cache
self.cache_k[:bsz].scatter_(dim=1, index=scatter_pos, src=xk[:, -self.sliding_window:])
self.cache_v[:bsz].scatter_(dim=1, index=scatter_pos, src=xv[:, -self.sliding_window:])

---

我相信多数读者读到这里会跟我有一样的疑问，只让Q与前面Window Size的KV计算Attention，不会影响最终的预测精度吗？因为我们知道当前生成的token是由前面所有token共同决定的。而且论文中并没有特别详细说明，且给出的示意图(下图) 也有些让人费解。

SWA确实限制了每个token的Q只能关注固定大小(Window Size)内的其他token，然而，信息通过网络的传播并不仅仅局限于Window Size的大小，它还设计多层Transformer之间的信息传递。

2 Mixtral 8x7B (MoE)模型

前文说过 Mixtral-8x7B就是Mistral 7B的MoE模型，除了上述Mistral 7B中的特性以外，Mixtral-8x7B还引入了MoE结构。MoE(Mixture-of-Experts) 其实也不是一个新技术，早在1991年就已经被Michael Jordan 和 Geoffrey Hinton所提出 Adaptive mixtures of local experts , 而且关于MoE的发展在深度学习界也从未停止过 (所谓经典永不过时说的便是如此)，相关的papers综述这里提供一个写的不错的Blog供大家参考一下：Mixture-of-Experts (MoE) 经典论文一览

这里简单的解释一下什么是MoE，简单点说就是我让一个网络模型结构有多条分支，每条分支代表一个Expert(专家)，每个Expert都有其擅长的领域，当具体任务来临时，可以通过一个门空位Gate来具体选择采用哪一个或者哪几个Experts进行计算，这样的好处就是让每个Expert更专注特定领域，降低了不同领域数据对权重学习的干扰。当然在训练MoE模型时也要注意各个Experts负载均衡，防止赢者通吃，达不到想要的目的。

具体到Mixtral 8x7B模型中，其MoE的结构示意图如下所示

可以发现，相对于Llama ，Mixtral 8x7B模型将FFN替换为MoE FFN，还是直接看代码

class MoeLayer(nn.Module):
    def __init__(self, experts: List[nn.Module], gate: nn.Module, moe_args: MoeArgs):
        super().__init__()
        assert len(experts) > 0
# 定义experts，就是一组(8个)Llama FFN，
# Llama FFN就是两个Linear + Silu + Linear
        self.experts = nn.ModuleList(experts)
# gate也是一个Linear，这个Linear weight的维度是[hidden_dim , num_experts]
        self.gate = gate
        self.args = moe_args
    def forward(self, inputs: torch.Tensor):
# 更改input shape [bst,seq_len,hidden-dim] -> [bst*seq_len,hidden-dim]
        inputs_squashed = inputs.view(-1, inputs.shape[-1])
# Gate Linear 将输入线性映射到num_experts
# 即[bst*seq_len,hidden-dim] -> [bst*seq_len,num_experts]
        gate_logits = self.gate(inputs_squashed)
# topk排序
# weights返回topk的值
# selected_experts 返回topk的index
        weights, selected_experts = torch.topk(
            gate_logits, self.args.num_experts_per_tok
        )
# 对每个weight做softmax，归一化
        weights = nn.functional.softmax(
            weights,
            dim=1,
            dtype=torch.float,
        ).type_as(inputs)
        results = torch.zeros_like(inputs_squashed)
        for i, expert in enumerate(self.experts):
# 根据selected_experts确定weight的行id和列id
            batch_idx, nth_expert = torch.where(selected_experts == i)
# 通过上述id选择对应的加权数据 以及执行对应的expert，并将结果加权求和
            results[batch_idx] += weights[batch_idx, nth_expert, None] * expert(
                inputs_squashed[batch_idx]
            )
        return results.view_as(inputs)

3 Llama-2 70B vs Mixtral 8x7B

文章的最后，我们再来对比一下Llama-2 70B 和 Mixtral 8x7B 的参数量以及浮点运算量(FLOPs)

• Params

  Model	Attention	FeedForward	Layers	Others	Total
  Llama-2 70B	8192* 8192 * 2+ 8192 * 1024 * 2 + 8192 = 151003136	8192 * 28672 * 3 + 8192=704651264	80	8192 * 32000 * 2+ 8192 = 524296192	68976648192=68.98B
  Mixtral 8x7B	4096 * 4096 * 2 + 4096 * 512 * 2 + 4096 = 37752832	4096 * 8 + 8 * (4096 * 14336 *3) + 4096 = 1409323008	32	4096 * 32000 * 2+ 4096 = 262148096	46568574976=46.57B

• FLOPs
  计算FLOPs，我们就都以输入为2048的单batch作为基准计算，并且我们只计算矩阵乘法相关的FLOPs作为整体网络FLOPs的估算，Norm层的计算先忽略

Model	Attention	FeedForward	Layers	Others	Total
Llama-2 70B	2 * 2048 * 8192 * 8192 * 2 + 2* 2048 * 8192 * 1024 * 2 + 64 * 2 *2048 * 128 * 2048 * 2 = 7.55914244 *10^{11}	3 * 8192 * 28672 * 2048 *2 =2.88621802 *10^{12}	80	2 * 8192 * 32000* 2048 * 2 = 2.14748365*10^{12}	2.93518065*10^{14}=293.5TFLOPs
Mixtral 8x7B	2 * 2048 * 4096 * 4096 *2 + 2 * 2048 * 4096 * 512 * 2 + 32 * 2 *2048 * 128 * 2048 * 2 = 2.23338299 * 10^{11}	2048 * 4096 * 8 * 2 +3 * 4096 * 14336 * 2048 * 2 * 2 = 1.44324323 * 10^{12}	32	2* 2048 * 4096 * 32000 * 2= 1.07374182* 10^{12}	5.44043508* 10^{13} = 54.4TFLOPs

三、推理优化原理与源码实现

Mixtral的代码在推理时用到的一些trick，具体为：

• Sliding Window Attention (SWA，滑动窗口Attention)
• Rolling Buffer Cache（也被称为Rotating Buffer Cache，即旋转式存储的KV cache）
• Long-context Chunking（长上下文场景下的chunking策略，配合前两者食用）

  一、LLM推理的两阶段

一个常规的LLM推理过程通常分为两个阶段：prefill和decode。

1.1 Prefill

预填充阶段。在这个阶段中，我们把整段prompt喂给模型做forward计算。如果采用KV cache技术，在这个阶段中我们会把prompt过后得到的保存在cache_k和cache_v中。这样在对后面的token计算attention时，我们就不需要对前面的token重复计算了$x_k$、$x_v$，可以帮助我们节省推理时间。

在上面的图例中，我们假设prompt中含有3个token，prefill阶段结束后，这三个token相关的KV值都被装进了cache。

1.2 Decode

生成response的阶段。在这个阶段中，我们根据prompt的prefill结果，一个token一个token地生成response。

同样，如果采用了KV cache，则每走完一个decode过程，我们就把对应response token的KV值存入cache中，以便能加速计算。例如对于图中的t4，它与cache中t0~t3的KV值计算完attention后，就把自己的KV值也装进cache中。对t6也是同理。

由于Decode阶段的是逐一生成token的，因此它不能像prefill阶段那样能做大段prompt的并行计算，所以在LLM推理过程中，Decode阶段的耗时一般是更大的。

二、Sliding Window Attention

2.1 原理

从第一部分的介绍中，我们应该能感受到一点：LLM推理中的KV cache加速法，是非常典型的用“空间换时间”的操作。随着seq_len变长，cache中存储的数据量也越来越大，对显存造成压力。

所以，我们自然而然想问：有什么办法能减缓cache的存储压力呢？

注意到，cache的存储压力之所以变大，是因为我们的Attention是causal decoder形式的，即每一个token，都要和它之前所有的token做Attention，所以cache中存储的数据量才和seq_len正相关。如果现在我们转换一下思路，假设每一个token只和包含其本身在内的前W个token做Attention，这样不就能把cache的容量维持在W吗？而从直觉上来说，这样的做法也有一定的道理：对当前token来说，距离越远的token，能提供的信息量往往越低，所以似乎没有必要浪费资源和这些远距离的token做Attention。

这种Attention思路的改进，就被称为Sliding Window Attention，其中W表示窗口长度。这也是Mixtral 7b 和Mixtral 8 * 7b采用的方法，我们通过作者论文中的一张图，更清晰地来看下它和传统Attention的区别，这里W=3。

2.2 为什么能用滑动窗口

虽然滑动窗口的策略看起来很不错，不过你一定有这样的疑惑：虽然距离越远的token涵盖的信息量可能越少，但不意味着它们对当前token一点用处都没有。在传统的Attention中，我们通过Attention score，或多或少给这些远距离的token一定的参与度；但是在Sliding Window Attention中，却直接杜绝了它们的参与，这真的合理吗？

为了回答这个问题，我们来看一个例子，在本例中W=4，num_layers = 4，num_tokens = 10。

我们从layer3最后一个位置的token（t9）看起：

• 对于layer3 t9，它是由layer2 t9做sliding window attention得来的。也就是layer3 t9能看到layer2 t6 ~ t9的信息
• 再来看layer2 t6，它能看到layer1 t3 ~ t6的信息。也就是说对于layer3 t9，它最远能看到layer1 t3这个位置。
• 以此类推，当我们来到layer0时，不难发现，对于layer3 t9，它最远能看到layer0 t0这个位置的信息。
  欸你发现了吗！对于**layer3 t9**，虽然在每一层它“最远”只能看到前置序列中部分token，但是只要模型够深，它一定能够在某一层看到所有的前置tokens。

如果你还觉得抽象，那么可以想想CNN技术中常谈的“感受野”。当你用一个固定大小的卷积窗口，对一张原始图片做若干次卷积，得到若干张特征图。越深的特征图，它的每一个像素点看到的原始图片的范围越广。类比到我们的滑动窗口Attention上，从layer0开始，每往上走一层，对应token的感受野就往前拓宽W。

所以，Silding Window Attention并非完全不利用窗口外的token信息，而是随着模型层数的增加，间接性地利用起窗口外的tokens。

三、Rolling Buffer Cache

3.1 原理

当我们使用滑动窗口后，KV Cache就不需要保存所有tokens的KV信息了，你可以将其视为一个固定容量（W）的cache，随着token index增加，我们来“滚动更新” KV Cache。

下图给出了Rolling Buffer Cache的运作流程:

在图例中，我们做推理时喂给模型一个batch_size = 3的batch，同时设W = 3。此时KV Cache的容量为(batch_size, W)。我们以第1条prompt This is an example of ...为例：

• 在i时刻，我们对an做attention，做完后将an的KV值更新进cache中
• 在 i + 1时刻，我们对example做attention，做完后将example的KV值更新进cache中。此时对于第1条prompt，它在KV cache中的存储空间已满。
• 在 i + 2时刻，我们对of做attention，由于此时KV cache已满，所以我们将of的KV值更新进KV cache的0号位置，替换掉原来This的KV值。再后面时刻的token也以此类推。
• 不难发现，prompt中第i个token在KV cache中的存储序号为：**i % W**

  3.2 “旋转”从何而来

如果你读过Mixtral的源码，你可能会记得，在源码中管Rolling Buffer Cache叫Rotary Buffer Cache。而“Rotary”这个词很值得我们关注：为什么叫“旋转”呢“

我们再回到3.1的图例中：

还是对于第一条数据，我们往上添两个单词，假设其为This is an example of my last...。现在来到了单词last上，我们需要对它计算Sliding Window Attention。

不难理解，在W=4的情况下，last的Attention和example of my last相关。现在我们把目光放到图中的KV Cache上：它的存储顺序似乎不太对，如果我们想对last做Attention，就要对当前KV Cache中存储的元素做一次“旋转”，将其转回正确的位置。

所以，Rotary的意思就是：通过某种规则，将Cache中的数据旋转回正确的位置，以便能正确做Attention。这个规则在Mixtral源码中用一个unrotate函数来定义。在后文我们会详细看这个函数的运作方式。

四、Chunking

我们回忆一下目前为止Mixtral为了加速模型推理做的操作：

• 使用KV Cache，加速Decode过程
• 使用Sliding Window Attention和Rolling Buffer Cache，降低KV Cache存储压力
  你可能已经发现，这些以“空间换时间”的优化，都是针对Decode过程的。那么对于Prefill过程，我们能做什么优化呢？

相比于更耗时的Decode阶段，Prefill有一个更加突出的问题：long-context。过长的prompt会给显存带来压力。一个符合直觉的解决办法是：把prompt切成若干chunk，每次只喂给模型1个chunk，更新1次KV Cache。这样我们虽然牺牲了一些Prefill计算的并行性（所有tokens一起计算），却能帮助我们节省显存压力（尤其是在采用sliding window attention的情况下，KV Cache的尺寸是固定的而不是随seq_len增长时）。

一般情况下，我们设chunk_size = cache_window = sliding_window = W，也就是chunk和cache的尺寸都和滑动窗口的尺寸保持一致，都设为W。对这个参数设置我们再说明下：一般满足cache_window = sliding_window，这个不难理解，因为cache中存的是attention感受野范围内的token。而chunk_size可以不等于这两者（源码中也提供了相关处理）。只是chunk_size和这两者相等时，无论是从计算逻辑还是空间利用率上，都是更好的选择（现在觉得抽象没关系，后文会提供具体的图例，大家可以感受下）。

好，现在我们来看一个chunking的图例（来自Mixtral论文），假设输入的prompt为The cat sat on the mat and saw the dog go to，同时chunk_size = cache_window = sliding_window = 4

假设我们现在来到第三块chunk，它包含的词为the dog go to。我们要对这个chunk中的每一个token计算滑动窗口Attention，同时把每个token的Xk, Xv值更新进KV Cache。

• 图中row方向表示Xq，即你可以把row方向the dog go to的每一个token，当成是这个token过Wq后的Xq值
• 图中col方向表示Xk, Xv，即你可以把col方向The cat sat on the mat and saw the dog go to的每一个token，当成是这个token过Wk，Wv后的Xk，Xv值，这些值存储在KV Cache中
• 图中整个0/1数据块表示mask矩阵。它表示row方向的Xq应该和col方向的哪些Xk，Xv值做attention。
  现在我们已基本能理解这张图的含义，不过还有一点很奇怪：在这个图下的Past, Cache, Current表示什么意思呢？

我们牢记一点：只有1个KV cache（也可以理解成只有1个用于存放Xk值的cache_k，和1个用于存放Xv值的cache_v）。当我们遍历到某个chunk时，我们取出当前的cache和这个chunk做attention计算，然后再把这个chunk相关的KV值按Rolling Buffer Cache的方式更新进这个cache中。

回到我们的例子上，现在我们位于第3块chunk上，此刻cache中存储的Xk, Xv值，即是上图中间块维护的the mat and saw，因此只有中间块的最底下被标上了“cache”，因为它才是此时真正的cache。而最左侧past块维护的则是前一个时刻的cache。最右侧的current块维护的the dog go to是即将被更新进cache的Xk, Xv值。这就是past, cache和current的含义。

注意到虽然图中画出了past块，但这并不意味着计算第3块时要把past块也取出（此时past块代表的cache早就被更新了）。论文中这样画只是更方便我们了解cache更新迭代和计算的过程。（悄悄吐槽下，虽然论文中的这些图画得很好很精练，但是少了很多关键信息的文字介绍，容易给人造成似懂非懂的感觉）

五、Chunking推理全流程图解

我们用图解的方式把整个推理流程串一遍，好知道代码在做一件什么事情

5.1 输入数据

假设推理时batch_size = 3，且有chunk_size = cache_size = sliding_window = 4，则这个batch的prompts可表示成下图（每个方块表示1个token，同色方块属于同个prompt）：

• 我们首先将chunk0送入模型，此时KV cache为空

• 对chunk中的每个token计算Xq，Xk，Xv，用于计算SWA（Sliding Window Attention）。图中刻画了计算时用到的mask矩阵。在Mixtral源码中使用Xformers库的相关API来完成Attention相关的计算（这个库的好处是加速Attention计算）。BlockDiagonalCausalMask（全称是BlockDiagonalCausalLocalAttentionMask）是这个库下提供的一种mask方法，它可以这样理解：
  Xformers官方文档在这一块的介绍不太全面，对初次使用Xformers的朋友其实不太友好，所以在这里我做了可视化，方便后续大家对代码的理解。

• chunk0的SWA计算完毕后，我们将每个token对应的Xk, Xv值存入cache。在源码中，我们会通过一个规则确定每个token的KV值在KV cache中的存储位置，这样也方便我们做unrotate操作（见本文3.2部分）时能把cache中存储的元素旋转回正确的位置。

• 最后，对于KV cache，它的position序号的排布顺序是从左至右，从上到下的，即：

  Cache position index:
  
  0 | 1 | 2  | 3
  4 | 5 | 6  | 7
  8 | 9 | 10 | 11

(2) chunk1

• 对于chunk1中维护的tokens，我们正常计算他们的xq，xk，xv。
• 取出当前KV Cache中存储的KV值，和chunk计算出来的KV值进行拼组，计算SWA（如图所示，mask矩阵的row行，每个色块由两部分组成：当前cache + 当前chunk）
• 在计算SWA的mask矩阵时，我们同样采用Xformers库，这时调用的是BlockDiagonalCausalLocalAttentionFromBottomRightMask类，和chunk0调用的BlockDiagonalCausalLocalAttentionMask相比，它的主要不同在“FromBottomRight”上，也就是对于每个block，它从右下角开始以窗口长度为W（本例中W=4）的形式设置mask矩阵。
• 计算完chunk1的SWA后，我们将chunk1的KV值更新进KV Cache中

  (3) chunk2

最后我们来看chunk2，这个chunk比较特殊，因为在这个chunk内，每一个prompt维护的序列长度是不一样的，3个prompt维护的tokens分别为[[8, 9, 10, 11], [8, 9], [8]]。

• 同样，我们计算chunk2的每个tokens的Xq，Xk，Xv
• 取出当前KV cache，与chunk2的相关结果做Attention计算，依然是采用Xformers的BlockDiagonalCausalLocalAttentionFromBottomRightMask类
• 把chunk2计算的KV结果更新进KV Cache。我们特别关注第2、3条prompt（绿红色块）更新后的KV cache结果。按照3.1中rolling buffer cache设置的放置方式，这两条prompt中KV值是非顺序存放的。例如对于第2条prompt，它KV值的存放顺序是[8, 9, 6, 7]。因此如果我们想继续对它做decode，就要把KV cache的值unrotate回[6, 7, 8, 9]，以此类推。

事实上，无论是prefill还是decode，无论是哪个chunk，只要涉及到用当前cache和chunk（在decode阶段则是token）做attention计算，我们都需要把cache中的KV值排布**unrotate**一遍。unrotate的结果就是，如果cache中的值已经是按顺序排布的，那就照常输出；如果是非顺序排布的，那就排好了再输出。由于在Mixtral源码中，这块数据处理逻辑比较复杂，又没有写注释，所以很多朋友读到unrotate的部分可能一头雾水。因此这里特地画出，帮助大家做源码解读。

一个新例子：chunk_size != W

在前文我们说过，一般设chunk_size = cache_window = sliding_window，我们也说过这个设置并不绝对，一般cache_window和sliding_window相等，但是chunk_size却不一定要和它们相等。

所以我们来看一个chunk_size和其余两者不等的例子。在这个例子中，chunk_size = 5, cache_window = sliding_window = 3

和5.2中的示例一样，对于每个chunk都主要分成三个阶段：更新前的KV Cache，SWA，更新后的KV cache。其中前两个阶段和5.2的示例差别不大，我们主要来关注下第三个阶段：更新KV Cache。

不难理解，对于每个chunk来说，只有倒数W个token的KV值才应该进KV cache。例如对prompt0的chunk0，我们自然而然会认为用它更新KV cache后，KV cache中token的排布应该是**[2, 3, 4]**，但真的是这样吗？

上图显示了prompt0的不同chunk更新KV cache后的结果，可以发现，chunk0更新KV cache后，元素的排布方式是**[3,4,2]（而不是我们认为的[2,3,4]）；chunk1更新KV cache后，元素的排布方式是[9, 7, 8]（而不是我们认为的[7, 8, 9]**）。这是因为整个更新过程严格遵循第三部分的Rolling Buffer Cache的更新原则（这样我们才能使用一套unrotate准则应对chunk_size等于和不等于cache_window/sliding_window的情况）。详细的更新过程已经在图例中画出。

同样，我们每次在使用KV Cache计算Attention时，也要注意用unrotate方法将KV Cache中的元素先按顺序排布好。

六、一些关于源码的hint

在写这篇文章时，本来是打算把源码一起讲的。但是写到这里发现，其实代码中最难理解的部分，在这篇文章中已经做了可视化了，剩下的代码细节对读者们来说应该没难度。在这里再给一些hint（应该也是读者最难理解的part）：

• 代码中的**RotatingBufferCache**类，用来定义一个KV cache。从始至终只有1个KV cache（或理解成1个cache_k + 1个cache_v），它在prefill和decode阶段不断被更新
• 代码中**CacheView**类，用来操作KV cache（正如它的命名一样，它是cache的视图）。如果说RotatingBufferCache用来管理cache的结构，那么CacheView则对cache中的具体数据进行更新、排序等操作。
• 代码中**RotatingCacheInputMetadata**类，用来定义如何生成当前chunk的KV cache信息。从上面的例子中我们知道，当前chunk计算出的KV值是要被更新进KV cache中的，那么chunk中的哪些token要被更新进KV cache中（例如chunk_size != sliding_window/cache_window时，只有倒数W个token要被更新进KV cache中）？这些token的KV值在cache中要存放在什么位置？诸如此类的信息，我们都在RotatingCacheInputMetadata中定义。
• 代码中**unrotate**方法，用来定义如何把KV cache中的元素正确排布，以便做Attention
• 代码中**interleave_list**方法，用来定义Attention mask矩阵中的col方向元素排布（例如5.2（2）中的中间部分的图）。interleave是“交织”的意思。什么是“交织”呢？就是prompt0 cache + prompt0 chunk + prompt 1 cache + prompt1 chunk + prompt2 cache + prompt2 chunk这样插入式交替排布的意思。