专家混合 (MoE)

相关源文件

• README.md
• model.py

目的与范围

本文档描述了 Grok-1 中实现的专家混合 (MoE) 架构。它涵盖了 MoE 系统的核心组件、路由机制、专家选择流程以及 MoE 如何与整体 Transformer 架构集成。有关更广泛的 Transformer 架构的信息，请参阅Transformer 架构。有关与 MoE 相关的内存管理详细信息，请参阅内存和 KV 缓存。

Grok-1 中的 MoE 简介

专家混合是一种条件计算技术，它将模型的参数划分为多个“专家”网络，其中只有一部分专家处理每个输入 token。这使得 Grok-1 能够扩展到 3140 亿参数，同时在推理过程中保持合理的计算要求。

Grok-1 采用以下规范实现 MoE：

• 每个 MoE 层有 8 个专家网络
• 每个 token 选择 2 个专家
• MoE 层集成到 Transformer 层的馈送前（feed-forward）组件中

来源：README.md24-30

MoE 架构概述

下图说明了 Grok-1 中 MoE 系统的总体架构

MoE 系统架构

来源：model.py1010-1102 model.py272-400

核心组件

Grok-1 中的 MoE 实现包含三个主要组件：

1. Router：确定哪些专家应处理每个 token。
2. Expert Networks：转换输入嵌入的馈送前神经网络。
3. MoE Layer：协调路由和专家计算。

组件关系

来源：model.py208-269 model.py272-400 model.py1010-1102 model.py420-486

路由机制

Router 负责确定哪些专家应处理每个 token。它为每个 token-专家对计算路由概率，并选择 top-k 个专家（在 Grok-1 中 k=2）。

Router 实现

Router 类实现：

1. 通过线性投影计算路由 logits。
2. 通过 softmax 将 logits 转换为概率。
3. 为每个 token 选择 top-k 个专家。

Router 使用简单的线性投影来计算路由 logits。

然后使用这些概率来选择哪些专家将处理每个 token，并选择 top-k（Grok-1 为 2）个专家。

来源：model.py208-269 model.py294-299

专家选择与计算

一旦 Router 确定了哪些专家应处理每个 token，MoE 层将执行以下步骤：

1. 专家选择：根据路由概率为每个 token 选择 top-2 专家。
2. Token 分配：将输入 token 路由到其指定的专家。
3. 专家计算：每个专家对其分配的 token 执行计算。
4. 输出组合：来自不同专家的输出将根据门控权重进行组合。

专家计算流程

Grok-1 中的每个专家都是一个馈送前网络，其结构类似于：

Linear_V + Linear → GELU → Linear_1

来源：model.py359-393 model.py977-1007

与 Transformer 架构集成

MoE 层取代了 Transformer 解码器层中的标准馈送前网络。在 Grok-1 架构中，是否使用 MoE 的决定由 num_experts 参数决定。

DecoderLayer 类中的实现显示了如何做出该决定。

来源：model.py1074-1093

性能考量

Grok-1 的 README 指出，MoE 层的实现并未针对效率进行优化。

本仓库中的 MoE 层实现效率不高。选择此实现是为了避免需要自定义内核来验证模型正确性。

此实现使用 JAX 的 shard_map 进行分布式计算，但该方法侧重于正确性而非性能优化。

MoE 实现包括对专家矩阵乘法的特殊处理，以支持分布式计算。

来源：README.md18-19 model.py319-337 model.py339-357

内存效率

Grok-1 的 MoE 架构允许模型扩展到 3140 亿参数，同时保持与小得多得密集模型相似的计算占用空间。这是因为：

1. 每个 token 在每个层中仅激活 8 个专家中的 2 个。
2. 与完全密集模型相比，这大约将每个 token 的计算量减少了 75%。
3. 模型可以拥有更多参数，而计算时间不会成比例增加。

正如 README 中提到的，Grok-1 的内存需求仍然很大。

由于模型规模庞大（3140 亿参数），需要具有足够 GPU 内存的机器才能使用示例代码测试模型。

来源：README.md18-19

量化支持

Grok-1 中的 MoE 实现支持 8 位量化以减少内存使用，这可以从 `_inference_call` 方法中对量化权重的条件处理中看出。

此实现使模型能够使用较低精度的权重表示运行，从而降低内存要求，同时保持合理的准确性。

来源：model.py359-393 model.py38-51

总结

Grok-1 中的专家混合（MoE）架构通过为每个 token 条件性地激活一小部分参数，使得模型能够扩展到 3140 亿参数。该实现包括：

1. 一个决定每个 token 使用哪些专家的路由器。
2. 一种为每个 token 选择 top-2 专家的机制。
3. 处理分配的 token 的专家网络。
4. 一个根据路由权重组合专家输出的系统。

虽然此实现并未针对性能进行优化，但它展示了 MoE 如何用于构建计算需求可控的超大规模语言模型。

来源：README.md24-30 model.py272-400