TensorRT-LLM中的专家并行性

专家混合模型 (MoE)

专家混合（MoE）架构最近被广泛使用，例如Mistral Mixtral 8x7B。具体来说，MOE的结构支持多个并行的前馈神经网络（FFN）层（称为专家）来替代密集模型中的单个FFN层。当令牌到达时，路由器层为每个令牌选择TopK专家。然后，令牌的相应隐藏状态分别被分发到所选的TopK专家。因此，每个专家都会接收到多个令牌的隐藏状态。

_{Switch Transformer中的MOE结构：https://arxiv.org/pdf/2101.03961.pdf}

张量并行与专家并行

如果MoE模型无法适应单个GPU的内存，那么多GPU上的并行性是必要的。我们支持了两种MoE结构的并行模式，张量并行（默认模式）、专家并行以及两者的混合模式。

张量并行均匀地分割每个专家的权重并将它们分配到不同的GPU上，这意味着每个GPU持有所有专家的部分权重，而专家并行则均匀地将一些专家的完整权重分配到不同的GPU上，这意味着每个GPU持有部分专家的完整权重。因此，张量并行组中的每个GPU等级接收所有专家的所有令牌的隐藏状态，然后使用部分权重进行计算，而对于专家并行，每个GPU等级仅接收该等级上专家的部分令牌的隐藏状态，然后使用完整权重进行计算。

当同时启用张量并行和专家并行时，每个GPU处理一部分专家权重矩阵（如在EP模式下），并且这些权重在多个GPU之间进一步切片（如在TP模式下）。这种混合方法旨在更均匀地平衡GPU之间的工作负载，提高效率并减少仅使用EP模式时可能出现的瓶颈。

如何启用

默认的并行模式是张量并行。您可以通过在调用convert_coneckpoint.py时设置--moe_tp_size和--moe_ep_size来启用专家并行或混合并行。如果仅提供--moe_tp_size，TRT-LLM将使用张量并行来处理MoE模型；如果仅提供--moe_ep_size，TRT-LLM将使用专家并行；如果两者都提供，则将使用混合并行。

确保moe_tp_size和moe_ep_size的乘积等于tp_size，因为所有GPU上的MoE并行总数必须与模型其他部分的并行总数相匹配。

与MoE结构相关的其他参数，例如num_experts_per_tok（在之前的上下文中称为TopK）和num_local_experts,可以在模型的配置文件中找到，例如Mixtral 8x7B模型的配置文件。 )