LLM Systems 17: System for MoE Models

0. Recap

  • Pipeline Parallelism
    • split by layers (horizonal split)
    • eliminate the bubbles (idle)
    • interleaving forward/backward
  • Tensor Parallelism
    • split the matrix computation over multiple GPUs
    • usually within a single node

1. Transformer Mixture-of-Expert Model

1.1 Motivation

问题：大模型训练的主要挑战来自计算量，如下图所示：


  • 对 Dense 模型进行扩展非常困难，因为 FLOPs 随着模型规模增加明显（考虑经典 Dense Transformer 模型，其 FLOPs 与层数 \(l\) 成线性关系，与隐藏层维度 \(h\) 成平方关系）。
  • Solution：Sparse model
    • 增加模型参数量与能力，而不需要显著增长计算量；
    • 典型代表：MoE。
  • MoE 优点：
    • 更易于扩展到大参数量模型；
    • 预训练比 Dense 模型快；
    • 对比相同参数量的 Dense 模型，MoE 模型推理也更快。

MoE 模型结构：

  • 使用以下模块替换 Transformer Layer 中的 FFN：
    • 多个小的专家，每个专家是一个神经网络（如 FFN）；
    • 一个 门控神经网络，用于根据输入 token 决定激活哪些专家；
  • MoE 不同于混合专家学习（Mixture-of-Expert learning），后者是一种学习算法，用于学习多个预测模型的均值权重。

1.2 Switch Transformer architecture

  • 每个 token 只激活一个 expert，如下图所示：


  • tokens （输入 shape：\([B, S, H]\)） 经过 router 后（权重 shape 为：\([H, E]\)），输出 tensor shape 为： \([B, S, E]\) ，其中 \(B, S, H, E\) 分别表示：batch size、sequence、hidden dim 和专家数量。选择输出的 E 个 logits（经过 Softmax 变为 probs） 中最大值对应专家，激活该专家对对应 token 进行 FFN 计算。FFN 输出乘以对应的 probs，得到最终的 MoE 层输出。
  • 门控网络 \(G\) 学习把 input 中的一个 token 发送给哪个 expert，具体过程形式化描述如下：
    • 计算原始得分 (Logits)： 输入的 token 特征 \(x\) 通过路由线性层（Router），计算出针对各个专家的原始得分（Logits）\(H(x) = x \cdot W_g\) ；
    • Top-K 筛选： 使用 \(\text{KeepTopK}(H(x), k)\) 筛选出得分最高的 \(k\) 个专家（在 Switch Transformer 中 \(k=1\)，即单专家路由），将未选中的专家得分设为 \(-\infty\) ；
    • Softmax 归一化： 通过 Softmax 计算门控概率/权重： \(G(x) = \text{Softmax}(\text{KeepTopK}(H(x), k))\) ；
    • 专家计算与加权： 最终输出 \(y\) 是 Top-k 专家输出 \(E_i(x)\) 与门控权重 \(G(x)_i\) 的乘积和： \(y = \sum_{i=1}^{k} G(x)_i E_i(x)\) 。

1.3 Shared-routed Experts

  • Motivation：所有 token 总是经过一个固定的 shared expert；
  • 不同专家作用：
    • Shared expert：calculating common knowledge；
    • Routed experts: calculated token-specific knowledge. First from Deepspeed-MoE. later in DeepSeek-MoE。
  • 计算流程如下图所示：


Expert 网络学到了什么？

  • Encoder experts tend to specialize in token groups or shallow concepts (e.g., punctuation, proper nouns).
  • Decoder experts exhibit less specialization.
  • In multilingual setups, experts do not specialize in specific languages due to token routing and load balancing.

Router 几何解释：

Router 本质可以理解为：高维特征空间中的线性超平面划分（finding linear boundaries among expert centroids）。

具体解释如下：

从数学公式看，Router 计算的是 \(H(x) = h \cdot W_g\)，其中 \(h\) 是输入 Token 的高维特征向量，而 \(W_g\) 是 Router 的权重矩阵。

从几何角度来看，这相当于：

  1. 专家的质心/方向向量（Expert Centroids）：
    • 矩阵 \(W_g\) 中的每一列 \(W_{g, i}\)，都可以看作是第 \(i\) 个专家在高维特征空间中的“质心（Centroid）”或标志性方向向量。
  2. 点积与相似度投影：
    • \(h \cdot W_{g, i}\) 的点积计算，几何上等于计算 Token 向量 \(h\) 在第 \(i\) 个专家方向上的投影长度（代表相似度得分）。
  3. KeepTop1 划分线性决策边界（Linear Boundaries）：
    • KeepTop1 挑选的是点积最大的专家，即 \(\arg\max_i (h \cdot W_{g, i})\)。
    • 对于任意两个专家 \(i\) 和 \(j\)，分割它们领地的分界线满足 \(h \cdot W_{g, i} = h \cdot W_{g, j}\)，即 \(h \cdot (W_{g, i} - W_{g, j}) = 0\)。这是一个穿过原点的高维线性超平面。
    • 3 个路由专家就相当于用线性超平面将整个高维 Token 特征空间切割成了 3 个凸多面体区域（类似 Voronoi 胞腔）。
      • Voronoi 胞腔：向量近似最近邻检索（Approximate Nearest Neighbor Search，ANNS）领域常见概念。
      • 可以用 ANNS 中概念理解 Router：Token Hidden State ＝ Query 向量，Router 权重 \(W_{g, i}\) ＝ 专家的聚类质心（Centroid），Router 计算 ＝ 最大内积搜索（MIPS）。
      • BTW：Self-Attention 计算，也可以理解为 MIPS 检索（\(QK^T\)）；外加 Softmax 归一化以及乘以 Value。

MoE 如何在 FLOPs 不变情况下扩展假设空间

从假设空间的几何结构来看，MoE 并不是用一个高维的全局连续函数去硬拟合所有数据，而是将庞大的假设空间拆解为由 Router 门控连接的 \(N\) 个局部子假设空间（Sub-hypotheses）。

  • 每一个输入的 Token，在运行时由 Router 动态投射到最匹配其特征的特定专家（子函数）中。

  • 这种机制的本质，是用一组 分段局部拟合函数 去刻画极其高维、非线性的语言数据流形。因为单 Token 实际激活的子空间计算量与 Dense 模型完全一致，所以它能以常数级别的 FLOPs，在全局假设空间中记忆并建模截然不同的模式（如代码、多语言、事实知识）。

  • 模型结构的创新，本质上是否可以理解为：在相同的约束下（FLOPs、Memory、Communication 等），提升假设空间表示的多样性/复杂性？

    • 模型结构的创新，本质上就是在物理约束（FLOPs、显存带宽、通信延迟、训练稳定性）构成的帕累托边界（Pareto Frontier）内，通过重构归纳偏置（Inductive Bias），实现假设空间有效表达容量（Effective Expressive Capacity）的最大化。

2. Training and Inference for MoE

2.1 Training of MoE: Expert Parallelism


  • 每个 worker/device 一个 Expert；
  • 其他网络结构在不同 device 间复制（DP）；
  • 需要快速的 All2All 通信：
    • 经过 Router/Gating 后，根据输出 probs 对 tokens 进行 All2All Dispatch；
    • 每个 device 上 Expert 进行计算；
    • 进行 All2All Combine 通信，将计算后结果发送给原本输入 token 所在 device；

Token 视角的 MoE forward 过程：


  • token 经过 Router 可能发送给不同的 Expert 进行计算；
  • Expert 计算完成后，发送给 token 输入所在 device；
  • 继续进行残差连接 / Norm 计算。

Load Balancing in MoE Training

  • 增加额外的损失函数：Expert-Level Balance Loss，以避免路由塌陷（routing collapse）导致的 token 全部被路由至少量 token，同时会引起负载不均，导致整体 GPU 吞吐浪费；
  • 专家级负载均衡损失函数定义：

\[ L_{\text{ExpBal}} = \alpha_1 \cdot M \cdot \sum_{i=1}^{M} f_i \cdot P_i \]

其中：

  • \(M\)：专家的总数量（#experts）。
  • \(\alpha_1\)：辅助损失的超参数权重（乘数系数），用于控制负载均衡损失在总 Training Loss 中的占比。
  • \(f_i\)（实际路由占比 / 频次）： \(f_i = \frac{\text{\#tokens routed to expert } i}{\text{\#tokens}}\) 代表在一个 Batch 中，真正被 Top-K 选中并分配给专家 \(i\) 的 Token 比例。
  • \(P_i\)（门控概率均值）： \(P_i = \frac{1}{\text{\#tokens}} \sum_{t=1}^{\text{\#tokens}} s_{i, t}\) 代表在一个 Batch 中，Router 的 Softmax 输出中分配给专家 \(i\) 的平均门控概率。

2.2 MoE Inference

  • MoE 推理性能决定因素：
    • 模型大小
    • 激活专家数
    • 显存带宽
  • 默认实现：
    • 所有专家都放在 GPU 显存

Optimizing MoE inference

  • 系统设计目标：
    • minimize the critical data path per device
    • maximize the achievable aggregate memory bandwidth
  • 具体实现：
    • Token 路由与分组聚合：将具有相同计算路径（目标专家）的 Token 组合在一起，减少单设备的数据重组开销。
    • 并行协同通信调度：针对不同并行策略重叠并优化 All-to-All 通信。
    • 算子/Kernel 级融合：优化 Transformer 和 MoE 相关算子以提升单设备计算效率。

Expert Parallelism and Tensor-Parallelism


  • 共 16 个 GPUs
    • Attention 层：TP=4，DP=4；
    • MoE 层：每个 Expert TP=2，EP=8；
      • Expert 0 权重切分放置在 GPU 0 和 GPU 8；
  • 完成 Attention TP 计算 + AllReduce 后，每个 DP Group 内的 4 个 GPU 上结果是完全相同的。如果直接进行 All2All Dispatch，会导致重复发送和带宽浪费。
    • Solution：将 TP 组内相同的 Token 集合按 Hidden 维度切开，或者只由 TP 组内的特定 GPU 负责分发。
  • 考虑 GPU 0 完成 Attention 计算后，其将经过 Router 后的 tokens 重排，把分配给 Expert 0 的数据的 \(\frac{H}{\text{\#Attention TP} \cdot \text{\#Expert TP}} = \frac{H}{4 \cdot 2}\) 维度发送给 Expert 0 所在的 GPU 0 和 GPU 8；
  • 进入 MoE FFN 之前，Expert TP Group 进行 AllGather 通信获取全量输入 \([B, S, H]\)；
  • 进行 Expert 内的 TP Forward 计算；
  • All2All Combine 发送计算结果到 token 所在 GPU；
    • 这里理论上也有通信优化：沿着 \(H\) 维度先切分做 All2All，然后 Attention TP Group 内做 AllGather 利用高速机内互联；

Optimizing MoE Kernels

  • MoE Specific Optimizations:
    • fuse the gating function into a single kernel
      • GEMM + Softmax + Top-K + Mask & Index
    • dense token-to-expert mapping table
  • Result: over 6x reduction in MoE Kernel related latency

Optimizing All-to-All Communication

  • Expert parallelism requires all-to-all communication between all expert parallel devices; the latency increases linearly with the increase in devices.
  • Optimization:
    • hierarchical all-to-all communication pattern: reduces the communication hops.
      • 机内预聚合（Intra-Node Aggregation）：在同一个服务器节点内部，几张 GPU 优先通过高速 NVLink，把要发往同一个远端目标节点的 Token 集中打包。
      • 机间网卡传输（Inter-Node Transfer）：利用节点间网卡，把打包好的大块数据包一次性跨节点发送给目标服务器（将 \(N \times N\) 的分散通信收拢为节点间的批量传输）。
      • 机内解包分发（Intra-Node Distribution）：目标服务器收到大包后，再通过本地 NVLink 分发给具体的专家 GPU。
    • parallelism-coordinated communication optimization: schedules communications based on the model's parallelism strategy to minimize overhead.

3. DeepSeek MoE (V3 model)

3.1 Model Structure


  • Fine-grained experts: each FFN is split to k smaller experts, total kN (N=original experts)
  • shared experts + routing experts
  • topk weighted average of routing experts (activating kM)

DeepSeek V3 MoE (670B)


参数名称 具体数值 架构设计含义
Model Dimension (\(H\)) 7168 隐层特征向量维度。
Total Layers 61 全模型共 61 层 Transformer Block。
Dense Layers 3 最底层的前 3 层为传统 Dense 层，不使用 MoE 路由，用于稳定早期特征。
FFN Inter Dim 18432 传统 Dense FFN 的中间层维度。
MoE Dim 2048 细粒度小专家的 FFN 中间维度（\(18432\) 被细分为了单专家 \(2048\) 维度）。
Shared Experts 1 个 1 个总是被激活的共享专家，捕获通用常识。
Routed Experts 256 个 全网提供高达 256 个细粒度路由专家供筛选2
Activated Experts 8 个 (\(K=8\)) 每次前向推导从 256 个路由专家中激活得分最高的前 8 个。
Expert Groups 8 组 将 256 个专家划分为 8 个专家组。
Limited Groups 4 组 组限制路由：限制 Token 只能在最多 4 个组内挑选 8 个专家，用于优化跨卡通信拓扑。

Load Balancing in DeepSeek MoE

  1. 专家级负载均衡损失（Expert-Level Balance Loss）
  • 核心公式：

\[ L_{\text{ExpBal}} = \alpha_1 \sum_{i=1}^{\#experts} f_i P_i \]

  • 参数含义：
    • \(f_i\)：被实际分发/路由给专家 \(i\) 的 Token 硬选择比例（Fraction of tokens dispatched to expert \(i\)），归一化公式为 \(f_i = \frac{\#experts}{\#activated\_experts} \cdot \frac{\#tokens \text{ to expert } i}{\#tokens}\)。
    • \(P_i\)：Router 对所有 Token 预测给专家 \(i\) 的 软概率/门控权重均值（Average routing probability for expert \(i\)），公式为 \(P_i = \frac{1}{\#tokens} \sum_{t=1}^{\#tokens} s_{i,t}\)。
  • 作用：在微观层面强迫 Router 将流量均匀分配给 256 个小专家中的每一个，**防止路由塌陷，确保每一个专家都能在训练中得到参数更新。
  1. 设备/组级负载均衡损失（Device-Level Balance Loss）
  • 核心公式：

\[ L_{\text{DevBal}} = \alpha_2 \sum_{j=1}^{\#groups} f_j P_j \]

  • 参数含义：
    • \(f_j\)：专家组（Group/Device）\(j\) 内所有专家的 \(f_i\) 平均值（代表分发给第 \(j\) 个设备/组的总体 Token 流量）。
    • \(P_j\)：专家组 \(j\) 内所有专家的 \(P_i\) 累加和（代表第 \(j\) 个设备/组收到的总体门控概率）。
  • 背景与作用：
    • 痛点：即使每个专家个体是均衡的，但如果落在某个 GPU 节点（或卡组）上的所有专家恰好被同时频繁激活，持有该节点的 GPU 算力就会被撑爆，产生“慢卡（Straggler Effect）”，导致全网等待。
    • 作用：在宏观物理设备层面强迫不同 GPU 设备/节点组（Groups）之间的总体计算负担保持均衡，消除卡间计算不对等，从而最大化分布式专家并行（EP）的计算吞吐。

DeepSeek Libraries to Accelerate MOE

3.2 code walkthrough

3.2.1 DeepSeek MoE

import & file variable

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
import math
from dataclasses import dataclass
from typing import Tuple, Optional, Literal

import torch
from torch import nn
import torch.nn.functional as F
import torch.distributed as dist

from kernel import act_quant, weight_dequant, fp8_gemm


world_size = 1
rank = 0
block_size = 128
gemm_impl: Literal["bf16", "fp8"] = "bf16"
attn_impl: Literal["naive", "absorb"] = "absorb"

ParallelEmbedding

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
class ParallelEmbedding(nn.Module):
    """
    Embedding layer with parallelism support across distributed processes.

    Args:
        vocab_size (int): Vocabulary size.
        dim (int): Embedding dimension.
    """
    def __init__(self, vocab_size: int, dim: int):
        super().__init__()
        self.vocab_size = vocab_size
        self.dim = dim
        assert vocab_size % world_size == 0, f"Vocabulary size must be divisible by world size (world_size={world_size})"
        self.part_vocab_size = (vocab_size // world_size)
        self.vocab_start_idx = rank * self.part_vocab_size
        self.vocab_end_idx = self.vocab_start_idx + self.part_vocab_size
        self.weight = nn.Parameter(torch.empty(self.part_vocab_size, self.dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass for parallel embedding layer.

        Args:
            x (torch.Tensor): Input tensor containing token indices.

        Returns:
            torch.Tensor: Embedded representations.

        Raises:
            ValueError: If `world_size` is not defined.
        """
        if world_size > 1:
            mask = (x < self.vocab_start_idx) | (x >= self.vocab_end_idx)
            x = x - self.vocab_start_idx
            x[mask] = 0
        y = F.embedding(x, self.weight)
        if world_size > 1:
            y[mask] = 0
            dist.all_reduce(y)
        return y
  • 按词表维度切分，fwd 时 AllReduce 通信。

Linear

  • 支持 FP8，代码略。

Column/Row Parallel Linear

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
class ColumnParallelLinear(Linear):
    """
    Linear layer with column parallelism, splitting output features across distributed processes.

    Args:
        in_features (int): Number of input features.
        out_features (int): Total number of output features.
        bias (bool): Whether to include a bias term. Defaults to False.
        dtype (optional): Data type for the layer. Defaults to `torch.bfloat16`.
    """
    def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype = None):
        assert out_features % world_size == 0, f"Output features must be divisible by world size (world_size={world_size})"
        self.part_out_features = out_features // world_size
        super().__init__(in_features, self.part_out_features, bias, dtype)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass for column parallel linear layer.

        Args:
            x (torch.Tensor): Input tensor.

        Returns:
            torch.Tensor: Transformed tensor with column-parallel computation.
        """
        y = linear(x, self.weight, self.bias)
        return y


class RowParallelLinear(Linear):
    """
    Linear layer with row parallelism, splitting input features across distributed processes.

    Args:
        in_features (int): Total number of input features.
        out_features (int): Number of output features.
        bias (bool): Whether to include a bias term. Defaults to False.
        dtype (optional): Data type for the layer. Defaults to `torch.bfloat16`.
    """
    def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype = None):
        assert in_features % world_size == 0, f"Input features must be divisible by world size (world_size={world_size})"
        self.part_in_features = in_features // world_size
        super().__init__(self.part_in_features, out_features, bias, dtype)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass for row parallel linear layer.

        Args:
            x (torch.Tensor): Input tensor.

        Returns:
            torch.Tensor: Transformed tensor with row-parallel computation.
        """
        y = linear(x, self.weight)
        if world_size > 1:
            dist.all_reduce(y)
        if self.bias is not None:
            y += self.bias
        return y
  • ColumnParallel
    • 直接进行本地 forward 计算，无通信。
  • RowParallel
    • 本地 GEMM 计算完成后进行 AllReduce。

MLP

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
class MLP(nn.Module):
    """
    Multi-Layer Perceptron (MLP) used as a feed-forward layer.

    Attributes:
        w1 (nn.Module): Linear layer for input-to-hidden transformation.
        w2 (nn.Module): Linear layer for hidden-to-output transformation.
        w3 (nn.Module): Additional linear layer for feature transformation.
    """
    def __init__(self, dim: int, inter_dim: int):
        """
        Initializes the MLP layer.

        Args:
            dim (int): Input and output dimensionality.
            inter_dim (int): Hidden layer dimensionality.
        """
        super().__init__()
        self.w1 = ColumnParallelLinear(dim, inter_dim)
        self.w2 = RowParallelLinear(inter_dim, dim)
        self.w3 = ColumnParallelLinear(dim, inter_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass for the MLP layer.

        Args:
            x (torch.Tensor): Input tensor.

        Returns:
            torch.Tensor: Output tensor after MLP computation.
        """
        return self.w2(F.silu(self.w1(x)) * self.w3(x))
  • SwiGLU（Swish Gated Linear Unit）
    • 两个 ColumnParallel + 一个 RowParallel

Gate

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
class Gate(nn.Module):
    """
    Gating mechanism for routing inputs in a mixture-of-experts (MoE) model.

    Attributes:
        dim (int): Dimensionality of input features.
        topk (int): Number of top experts activated for each input.
        n_groups (int): Number of groups for routing.
        topk_groups (int): Number of groups to route inputs to.
        score_func (str): Scoring function ('softmax' or 'sigmoid').
        route_scale (float): Scaling factor for routing weights.
        weight (torch.nn.Parameter): Learnable weights for the gate.
        bias (Optional[torch.nn.Parameter]): Optional bias term for the gate.
    """
    def __init__(self, args: ModelArgs):
        """
        Initializes the Gate module.

        Args:
            args (ModelArgs): Model arguments containing gating parameters.
        """
        super().__init__()
        self.dim = args.dim
        self.topk = args.n_activated_experts
        self.n_groups = args.n_expert_groups
        self.topk_groups = args.n_limited_groups
        self.score_func = args.score_func
        self.route_scale = args.route_scale
        self.weight = nn.Parameter(torch.empty(args.n_routed_experts, args.dim))
        self.bias = nn.Parameter(torch.empty(args.n_routed_experts, dtype=torch.float32)) if self.dim == 7168 else None

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Forward pass for the gating mechanism.

        Args:
            x (torch.Tensor): Input tensor.

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: Routing weights and selected expert indices.
        """
        scores = linear(x, self.weight)
        if self.score_func == "softmax":
            scores = scores.softmax(dim=-1, dtype=torch.float32)
        else:
            scores = scores.sigmoid()
        original_scores = scores
        if self.bias is not None:
            scores = scores + self.bias
        if self.n_groups > 1:
            scores = scores.view(x.size(0), self.n_groups, -1)
            if self.bias is None:
                group_scores = scores.amax(dim=-1)
            else:
                group_scores = scores.topk(2, dim=-1)[0].sum(dim=-1)
            indices = group_scores.topk(self.topk_groups, dim=-1)[1]
            mask = scores.new_ones(x.size(0), self.n_groups, dtype=bool).scatter_(1, indices, False)
            scores = scores.masked_fill_(mask.unsqueeze(-1), float("-inf")).flatten(1)
        indices = torch.topk(scores, self.topk, dim=-1)[1]
        weights = original_scores.gather(1, indices)
        if self.score_func == "sigmoid":
            weights /= weights.sum(dim=-1, keepdim=True)
        weights *= self.route_scale
        return weights.type_as(x), indices
  • 主要流程：
    • GEMM + Softmax 计算 probs
    • 组限制路由
    • 全局 Top-K 提取与权重归一化
    • weights 缩放

Expert

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
class Expert(nn.Module):
    """
    Expert layer for Mixture-of-Experts (MoE) models.

    Attributes:
        w1 (nn.Module): Linear layer for input-to-hidden transformation.
        w2 (nn.Module): Linear layer for hidden-to-output transformation.
        w3 (nn.Module): Additional linear layer for feature transformation.
    """
    def __init__(self, dim: int, inter_dim: int):
        """
        Initializes the Expert layer.

        Args:
            dim (int): Input and output dimensionality.
            inter_dim (int): Hidden layer dimensionality.
        """
        super().__init__()
        self.w1 = Linear(dim, inter_dim)
        self.w2 = Linear(inter_dim, dim)
        self.w3 = Linear(dim, inter_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass for the Expert layer.

        Args:
            x (torch.Tensor): Input tensor.

        Returns:
            torch.Tensor: Output tensor after expert computation.
        """
        return self.w2(F.silu(self.w1(x)) * self.w3(x))
  • 类似 MLP，但是 Linear 层无并行

MoE

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
class MoE(nn.Module):
    """
    Mixture-of-Experts (MoE) module.

    Attributes:
        dim (int): Dimensionality of input features.
        n_routed_experts (int): Total number of experts in the model.
        n_local_experts (int): Number of experts handled locally in distributed systems.
        n_activated_experts (int): Number of experts activated for each input.
        gate (nn.Module): Gating mechanism to route inputs to experts.
        experts (nn.ModuleList): List of expert modules.
        shared_experts (nn.Module): Shared experts applied to all inputs.
    """
    def __init__(self, args: ModelArgs):
        """
        Initializes the MoE module.

        Args:
            args (ModelArgs): Model arguments containing MoE parameters.
        """
        super().__init__()
        self.dim = args.dim
        assert args.n_routed_experts % world_size == 0, f"Number of experts must be divisible by world size (world_size={world_size})"
        self.n_routed_experts = args.n_routed_experts
        self.n_local_experts = args.n_routed_experts // world_size
        self.n_activated_experts = args.n_activated_experts
        self.experts_start_idx = rank * self.n_local_experts
        self.experts_end_idx = self.experts_start_idx + self.n_local_experts
        # router
        self.gate = Gate(args)
        # routed experts，只创建部分，EP
        self.experts = nn.ModuleList([Expert(args.dim, args.moe_inter_dim) if self.experts_start_idx <= i < self.experts_end_idx else None
                                      for i in range(self.n_routed_experts)])
        # shared expert
        self.shared_experts = MLP(args.dim, args.n_shared_experts * args.moe_inter_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass for the MoE module.

        Args:
            x (torch.Tensor): Input tensor.

        Returns:
            torch.Tensor: Output tensor after expert routing and computation.
        """
        shape = x.size()
        x = x.view(-1, self.dim)
        # 经过 router 获取 weights & indices
        weights, indices = self.gate(x)
        y = torch.zeros_like(x)
        counts = torch.bincount(indices.flatten(), minlength=self.n_routed_experts).tolist()
        for i in range(self.experts_start_idx, self.experts_end_idx):
            if counts[i] == 0:
                continue
            expert = self.experts[i]
            idx, top = torch.where(indices == i)
            # sum(weights & expert_1(x))
            y[idx] += expert(x[idx]) * weights[idx, top, None]
        z = self.shared_experts(x)
        if world_size > 1:
            dist.all_reduce(y)
        return (y + z).view(shape)
  • 见代码注释
  • 不是 All2All 通信实现，而是朴素的 local 计算 + 0 占位 + AllReduce（类似 ParallelEmbedding）

TransformerBlock

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
class Block(nn.Module):
    """
    Transformer block combining attention and feed-forward layers.

    Attributes:
        attn (nn.Module): Attention layer (MLA).
        ffn (nn.Module): Feed-forward network (MLP or MoE).
        attn_norm (nn.Module): Layer normalization for attention.
        ffn_norm (nn.Module): Layer normalization for feed-forward network.
    """
    def __init__(self, layer_id: int, args: ModelArgs):
        """
        Initializes the Transformer block.

        Args:
            layer_id (int): Layer index in the transformer.
            args (ModelArgs): Model arguments containing block parameters.
        """
        super().__init__()
        self.attn = MLA(args)
        self.ffn = MLP(args.dim, args.inter_dim) if layer_id < args.n_dense_layers else MoE(args)
        self.attn_norm = RMSNorm(args.dim)
        self.ffn_norm = RMSNorm(args.dim)

    def forward(self, x: torch.Tensor, start_pos: int, freqs_cis: torch.Tensor, mask: Optional[torch.Tensor]) -> torch.Tensor:
        """
        Forward pass for the Transformer block.

        Args:
            x (torch.Tensor): Input tensor.
            start_pos (int): Starting position in the sequence.
            freqs_cis (torch.Tensor): Precomputed complex exponential values for rotary embeddings.
            mask (Optional[torch.Tensor]): Mask tensor to exclude certain positions from attention.

        Returns:
            torch.Tensor: Output tensor after block computation.
        """
        x = x + self.attn(self.attn_norm(x), start_pos, freqs_cis, mask)
        x = x + self.ffn(self.ffn_norm(x))
        return x
  • Attention：MLA
  • FFN：MLP 或者 MoE
    • 前几层 Dense MLP；后续层 MoE。

Transformer

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
class Transformer(nn.Module):
    """
    Transformer model with positional embeddings, multiple layers, and output projection.

    Attributes:
        max_seq_len (int): Maximum sequence length for the transformer.
        embed (nn.Module): Embedding layer for input tokens.
        layers (torch.nn.ModuleList): List of transformer blocks.
        norm (nn.Module): Layer normalization applied after all blocks.
        head (nn.Module): Output projection layer mapping to vocabulary size.
        freqs_cis (torch.Tensor): Precomputed complex exponential values for rotary embeddings.
    """
    def __init__(self, args: ModelArgs):
        """
        Initializes the Transformer model.

        Args:
            args (ModelArgs): Model arguments containing transformer parameters.
        """
        global world_size, rank
        world_size = dist.get_world_size() if dist.is_initialized() else 1
        rank = dist.get_rank() if dist.is_initialized() else 0
        Linear.dtype = torch.float8_e4m3fn if args.dtype == "fp8" else torch.bfloat16
        Linear.scale_fmt = args.scale_fmt
        super().__init__()
        self.max_seq_len = args.max_seq_len
        self.embed = ParallelEmbedding(args.vocab_size, args.dim)
        self.layers = torch.nn.ModuleList()
        for layer_id in range(args.n_layers):
            self.layers.append(Block(layer_id, args))
        self.norm = RMSNorm(args.dim)
        self.head = ColumnParallelLinear(args.dim, args.vocab_size, dtype=torch.get_default_dtype())
        self.register_buffer("freqs_cis", precompute_freqs_cis(args), persistent=False)

    @torch.inference_mode()
    def forward(self, tokens: torch.Tensor, start_pos: int = 0):
        """
        Forward pass for the Transformer model.

        Args:
            tokens (torch.Tensor): Input tensor of token IDs with shape (batch_size, seq_len).
            start_pos (int, optional): Starting position in the sequence for rotary embeddings. Defaults to 0.

        Returns:
            torch.Tensor: Logits tensor of shape (batch_size, vocab_size).
        """
        seqlen = tokens.size(1)
        h = self.embed(tokens)
        freqs_cis = self.freqs_cis[start_pos:start_pos+seqlen]
        mask = None
        if seqlen > 1:
            mask = torch.full((seqlen, seqlen), float("-inf"), device=tokens.device).triu_(1)
        for layer in self.layers:
            h = layer(h, start_pos, freqs_cis, mask)
        h = self.norm(h)[:, -1]
        logits = self.head(h)
        if world_size > 1:
            all_logits = [torch.empty_like(logits) for _ in range(world_size)]
            dist.all_gather(all_logits, logits)
            logits = torch.cat(all_logits, dim=-1)
        return logits
  • Transformer
    • Embedding
    • TransformerBlock
    • LMHead + RMSNorm

3.2.2 DeepSpeed MoE

MoE

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team

from typing import Optional, Tuple

import torch
from torch import nn
from torch.nn import functional as F

from deepspeed.utils import groups, log_dist
from .experts import Experts
from .sharded_moe import MOELayer, TopKGate


class MoE(nn.Module):
    """Initialize an MoE layer.

    Arguments:
        hidden_size (int): the hidden dimension of the model, importantly this is also the input and output dimension.
        expert (nn.Module): the torch module that defines the expert (e.g., MLP, torch.linear).
        num_experts (int, optional): default=1, the total number of experts per layer.
        ep_size (int, optional): default=1, number of ranks in the expert parallel world or group.
        k (int, optional): default=1, top-k gating value, only supports k=1 or k=2.
        capacity_factor (float, optional): default=1.0, the capacity of the expert at training time.
        eval_capacity_factor (float, optional): default=1.0, the capacity of the expert at eval time.
        min_capacity (int, optional): default=4, the minimum capacity per expert regardless of the capacity_factor.
        use_residual (bool, optional): default=False, make this MoE layer a Residual MoE (https://arxiv.org/abs/2201.05596) layer.
        noisy_gate_policy (str, optional): default=None, noisy gate policy, valid options are 'Jitter', 'RSample' or 'None'.
        drop_tokens (bool, optional): default=True, whether to drop tokens - (setting to False is equivalent to infinite capacity).
        use_rts (bool, optional): default=True, whether to use Random Token Selection.
        use_tutel (bool, optional): default=False, whether to use Tutel optimizations (if installed).
        enable_expert_tensor_parallelism (bool, optional): default=False, whether to use tensor parallelism for experts
        top2_2nd_expert_sampling (bool, optional): default=True, whether to perform sampling for 2nd expert
    """

    def __init__(self,
                 hidden_size: int,
                 expert: nn.Module,
                 num_experts: int = 1,
                 ep_size: int = 1,
                 k: int = 1,
                 capacity_factor: float = 1.0,
                 eval_capacity_factor: float = 1.0,
                 min_capacity: int = 4,
                 use_residual: bool = False,
                 noisy_gate_policy: Optional[str] = None,
                 drop_tokens: bool = True,
                 use_rts: bool = True,
                 use_tutel: bool = False,
                 enable_expert_tensor_parallelism: bool = False,
                 top2_2nd_expert_sampling: bool = True) -> None:

        super(MoE, self).__init__()

        self.use_residual = use_residual
        self.enable_expert_tensor_parallelism = enable_expert_tensor_parallelism
        assert num_experts % ep_size == 0, f"Number of experts ({num_experts}) should be divisible by expert parallel size ({ep_size})"
        self.ep_size = ep_size
        self.expert_group_name = f"ep_size_{self.ep_size}"
        self.num_experts = num_experts
        self.num_local_experts = num_experts // self.ep_size

        log_dist(
            f'Creating MoE layer with num_experts: {num_experts} | num_local_experts: {self.num_local_experts} | expert_parallel_size: {self.ep_size}',
            [0])

        assert noisy_gate_policy is None or noisy_gate_policy in ['None', 'Jitter', 'RSample'], \
            'Unsupported noisy_gate_policy: ' + noisy_gate_policy
		# experts
        experts = Experts(expert, self.num_local_experts, self.expert_group_name)
        # TopKGate + MOELayer
        self.deepspeed_moe = MOELayer(TopKGate(hidden_size, num_experts, k, capacity_factor, eval_capacity_factor,
                                               min_capacity, noisy_gate_policy, drop_tokens, use_rts, None,
                                               top2_2nd_expert_sampling),
                                      experts,
                                      self.expert_group_name,
                                      self.ep_size,
                                      self.num_local_experts,
                                      use_tutel=use_tutel)
        if self.use_residual:
            self.mlp = expert
            # coefficient is used for weighted sum of the output of expert and mlp
            self.coefficient = nn.Linear(hidden_size, 2)

    def set_deepspeed_parallelism(self, use_data_before_expert_parallel_: bool = False) -> None:
        self._create_process_groups(use_data_before_expert_parallel_=use_data_before_expert_parallel_)

    def _create_process_groups(self, use_data_before_expert_parallel_: bool = False) -> None:
        # Create process group for a layer if needed
        if self.expert_group_name not in groups._get_expert_parallel_group_dict():
            print(f"No existing process group found, creating a new group named: {self.expert_group_name}")
            if (groups.mpu is None) or (not self.enable_expert_tensor_parallelism):
                # Condition 1 - no groups.mpu means no tensor parallelism
                # Condition 2 - disabling expert tensor parallelism on purpose
                groups._create_expert_and_data_parallel(
                    self.ep_size, use_data_before_expert_parallel_=use_data_before_expert_parallel_)
            else:
                # expert tensor parallelism is enabled
                groups._create_expert_data_and_model_parallel(
                    self.ep_size, mpu=groups.mpu, use_data_before_expert_parallel_=use_data_before_expert_parallel_)
        # Set the group handle for the MOELayer (deepspeed_moe) object
        self.deepspeed_moe._set_ep_group(groups._get_expert_parallel_group(self.expert_group_name))

    def forward(self,
                hidden_states: torch.Tensor,
                used_token: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """ MoE forward

        Arguments:
            hidden_states (Tensor): input to the layer
            used_token (Tensor, optional): default: None, mask only used tokens

        Returns:
            A tuple including output, gate loss, and expert count.

            * output (Tensor): output of the model

            * l_aux (Tensor): gate loss value

            * exp_counts (Tensor): expert count
        """
        output = self.deepspeed_moe(hidden_states, used_token)
        if self.use_residual:
            # Residual MoE
            output_mlp = self.mlp(hidden_states)
            if isinstance(output_mlp, tuple):
                output_mlp = output_mlp[0]  # Ignore the bias term for now
            coef = self.coefficient(hidden_states)
            coef = F.softmax(coef, dim=-1)
            output = output * coef[..., 0:1] + output_mlp * coef[..., 1:]
        return output, self.deepspeed_moe.l_aux, self.deepspeed_moe.exp_counts

Experts

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
class Experts(nn.Module):

    def __init__(self, expert: nn.Module, num_local_experts: int = 1, expert_group_name: Optional[str] = None) -> None:
        super(Experts, self).__init__()

		# 深拷贝，复制得到多个 experts
        self.deepspeed_experts = nn.ModuleList([copy.deepcopy(expert) for _ in range(num_local_experts)])
        self.num_local_experts = num_local_experts

        # TODO: revisit allreduce for moe.gate...
        for expert in self.deepspeed_experts:
            # TODO: Create param groups to handle expert + data case (e.g. param.group = moe_group)
            for param in expert.parameters():
                param.allreduce = False
                param.group_name = expert_group_name

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        chunks = inputs.chunk(self.num_local_experts, dim=1)
        expert_outputs: List[torch.Tensor] = []

        for chunk, expert in zip(chunks, self.deepspeed_experts):
            out = expert(chunk)
            if isinstance(out, tuple):
                out = out[0]  # Ignore the bias term for now
            expert_outputs += [out]

        return torch.cat(expert_outputs, dim=1)

MoeLayer

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
# All2All 通信实现

# Based on https://github.com/pytorch/pytorch/pull/40762
class _AllToAll(torch.autograd.Function):

    @staticmethod
    def forward(ctx: Any, group: dist.ProcessGroup, input: Tensor) -> Tensor:  # type: ignore
        ctx.group = group
        input = input.contiguous()
        output = torch.empty_like(input)
        dist.all_to_all_single(output, input, group=group)
        return output

    @staticmethod
    def backward(ctx: Any, *grad_output: Tensor) -> Tuple[None, Tensor]:
        return (None, _AllToAll.apply(ctx.group, *grad_output))


class MOELayer(Base):
    """MOELayer module which implements MixtureOfExperts as described in Gshard_.
    ::

        gate = TopKGate(model_dim, num_experts)
        moe = MOELayer(gate, expert)
        output = moe(input)
        l_aux = moe.l_aux

    .. Gshard_: https://arxiv.org/pdf/2006.16668.pdf

    Args:
        gate (torch.nn.Module):
            gate network
        experts (torch.nn.Module):
            expert network
    """

    def __init__(self,
                 gate: Module,
                 experts: Module,
                 ep_group_name,
                 ep_size,
                 num_local_experts: int,
                 use_tutel: bool = False) -> None:
        super().__init__()
        self.gate = gate
        self.experts = experts
        self.ep_group = None
        self.ep_size = ep_size
        self.ep_group_name = ep_group_name
        self.num_local_experts = num_local_experts
        self.time_falltoall = 0.0
        self.time_salltoall = 0.0
        self.time_moe = 0.0
        self.timers = SynchronizedWallClockTimer()
        self.wall_clock_breakdown = False

        self.use_tutel = use_tutel and TUTEL_INSTALLED

        if self.use_tutel:
            logger.info('Using Tutel optimizations.')
        elif use_tutel and not TUTEL_INSTALLED:
            logger.warning("Tutel optimization requested but not installed. "
                           "Proceeding without Tutel.")

    def _set_ep_group(self, ep_group):
        self.ep_group = ep_group
        self.gate._set_ep_group(ep_group)

    def forward(self, *input: Tensor, **kwargs: Any) -> Tensor:

        if self.wall_clock_breakdown:
            self.timers(MOE_TIMER).start()

        # Implement Algorithm 2 from GShard paper.
        d_model = input[0].shape[-1]

        # Initial implementation -> Reshape into S tokens by dropping sequence dimension.
        # Reshape into G groups so that each group can distribute tokens equally
        # group_size = kwargs['group_size'] if 'group_size' in kwargs.keys() else 1
        reshaped_input = input[0].reshape(-1, d_model)

        self.l_aux, C, E, indices_, locations_, gates_, self.exp_counts = self.gate(
            reshaped_input, input[1], True, self.use_tutel)
        S, M = reshaped_input.size(0), reshaped_input.size(1)
        # Resolve the capacity tensor once; every later use of C would otherwise force
        # its own device-to-host sync.
        C = int(C)

        if self.use_tutel:
            if not hasattr(self, '_tutel_dispatcher'):
                self._tutel_dispatcher = tutel_moe.fast_dispatcher(E, C, M, dispatch_dtype=reshaped_input.dtype)
            self._tutel_dispatcher.update(list(indices_.unbind(0)),
                                          list(locations_.unbind(0)),
                                          list(gates_.unbind(0)),
                                          capacity=C)
            # encode() returns a flat [e * c, m]; reshape to [e, c, m] so the tensor-parallel
            # drop_tokens/gather_tokens below split the capacity dim rather than the model dim.
            dispatched_input = self._tutel_dispatcher.encode(reshaped_input).view(E, C, M)
        else:
            slots = _route_slots(indices_, locations_, E, C)
            dispatched_input = _sparse_encode(reshaped_input, slots, E, C)

        if self.wall_clock_breakdown:
            self.timers(FIRST_ALLTOALL_TIMER).start()

        tensor_model_world_size = bwc_tensor_model_parallel_world_size(groups.mpu)
        if tensor_model_world_size > 1:
            # If the non-expert is tensor-parallel,
            # Whether expert is tensor-parallel or not , it will create
            # duplicate tokens on the tensor-parallel ranks.
            # drop duplicate tokens also doubles up as a communication
            # optimization as we are reducing the all-to-all communication volume.
            # 1: for not tensor-parallel expert,drop duplicate tokens to ensure
            # both correctness and reduce all-to-all communication.
            # 2: for tensor-parallel expert,drop duplicate tokens to reduce all-to-all
            # communication volume,before expert execution, it is necessary to perform
            # an allgather to ensure correctness,
            dispatched_input = drop_tokens(dispatched_input, dim=1)

        if self.ep_size > 1:
            dispatched_input = _AllToAll.apply(self.ep_group, dispatched_input)
        else:
            dispatched_input = dispatched_input.contiguous()

        if self.wall_clock_breakdown:
            self.timers(FIRST_ALLTOALL_TIMER).stop()
            self.time_falltoall = self.timers(FIRST_ALLTOALL_TIMER).elapsed(reset=False)

        if tensor_model_world_size > 1 and groups._get_expert_model_parallel_world_size() > 1:
            # if both expert and non-expert are tensor-parallel
            # the dropped duplicate tokens need to be gathered on each
            # tensor parallel rank again to ensure correctness
            dispatched_input = gather_tokens(dispatched_input, dim=1)

        # Re-shape after all-to-all: ecm -> gecm
        dispatched_input = dispatched_input.reshape(self.ep_size, self.num_local_experts, -1, d_model)
        expert_output = self.experts(dispatched_input)
        # Re-shape before drop_tokens: gecm -> ecm
        expert_output = expert_output.reshape(self.ep_size * self.num_local_experts, -1, d_model)
        if tensor_model_world_size > 1 and groups._get_expert_model_parallel_world_size() > 1:
            # if both expert and non-expert are tensor-parallel
            # drop duplicate tokens to ensure both correctness
            # and reduce all-to-all communication.
            expert_output = drop_tokens(expert_output, dim=1)

        if self.wall_clock_breakdown:
            self.timers(SECOND_ALLTOALL_TIMER).start()

        if self.ep_size > 1:
            expert_output = _AllToAll.apply(self.ep_group, expert_output)
        else:
            expert_output = expert_output.contiguous()

        if self.wall_clock_breakdown:
            self.timers(SECOND_ALLTOALL_TIMER).stop()
            self.time_salltoall = self.timers(SECOND_ALLTOALL_TIMER).elapsed(reset=False)

        if tensor_model_world_size > 1:
            # the dropped duplicate tokens need to be gathered on each
            # tensor parallel rank again for the tensor-parallel
            # non-expert of the next layer.
            expert_output = gather_tokens(expert_output, dim=1)

        if self.use_tutel:
            combined_output = self._tutel_dispatcher.decode(expert_output.view(E * C, M))
        else:
            combined_output = _sparse_decode(expert_output.view(E * C, M), slots, gates_, S)

        a = combined_output.reshape(input[0].shape)

        if self.wall_clock_breakdown:
            self.timers(MOE_TIMER).stop()
            self.time_moe = self.timers(MOE_TIMER).elapsed(reset=False)

        return a
  • MoeLayer forward 的核心链路如下：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
输入 Tensor [S, M]
  │
  ├──> 1. Gate 路由计算：得到 (indices, locations, gates, C)
  │
  ├──> 2. Dispatch 打包：调用 _sparse_encode（或 Tutel 加速）得到 [E, C, M] 格式
  │
  ├──> 3. TP 去重优化：当存在 Tensor Parallelism 时，调用 drop_tokens 丢弃重复 Token，削减通信量
  │
  ├──> 4. 第一次 All-to-All（Dispatch 通信）：调用 _AllToAll.apply 将 Token 发往对应 GPU
  │
  ├──> 5. 本地专家计算：调用 self.experts(dispatched_input) 进行 GEMM 矩阵乘法
  │
  ├──> 6. 第二次 All-to-All（Combine 通信）：调用 _AllToAll.apply 将计算结果发回原始 GPU
  │
  └──> 7. Combine 解包：调用 _sparse_decode，还原 Token 顺序并按门控权重加权求和，输出 [S, M]

4. Conclusion

  • LLM Mixture-of-Expert Model
    • Instead of a single dense FFN, using multiple FFNs (experts)
    • Routing network to select one/multiple experts
    • Shared-routed experts (DeepSpeed-MOE, DeepSeek MOE)
    • a few dense layers, then MOE (DeepSeek MOE)
  • Scalable training/inference
    • expert parallelism: split experts and replicate non-expert (GShard)
    • all-to-all communication for expert output
    • load balancing: grouping and avoid collapse (DeepSeek)
    • optimized kernel for MoE
ML System
#CMU 11868 #LLM System #Distributed Training #MoE
LLM Systems 17: System for MoE Models
https://arcsin2.cloud/posts/2026/10/4071970685/
作者
arcsin2
发布于
2026年10月31日
许可协议