TensorRT-LLM 中的专家并行#

专家混合（MoE）#

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

moe_structure

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

张量并行 vs 专家并行#

如果单个 GPU 的内存无法容纳 MoE 模型，则在多 GPU 上进行并行是必要的。我们支持 MoE 结构的两种并行模式，张量并行（默认模式），专家并行以及两者的混合。

tensor parallel vs expert parallel

张量并行均匀地分割每个专家的权重并将它们分配给不同的 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 模型 的配置文件。）