MoEBlaze: Breaking the Memory Wall for Efficient MoE Training on Modern GPUs

A1 主要贡献

本文旨在解决现代大规模混合专家（MoE）架构中日益严重的“内存墙”瓶颈。MoE架构固有的稀疏性虽降低了计算成本，但也引入了巨大的激活内存开销，这主要源于大型的令牌（token）路由缓冲区以及中间张量的实例化和缓冲需求。这种内存压力限制了GPU上可容纳的最大批量大小（batch size）和序列长度，并导致过多的数据移动，从而阻碍了性能和模型的有效扩展。

为应对这些挑战，本文提出了MoEBlaze，一个高内存效率的MoE训练框架。MoEBlaze通过系统协同设计的方法来解决上述问题，其核心贡献如下：

A3 背景知识与动机

本章节回顾了MoE训练的背景，并指出了当前限制其性能的关键系统瓶颈。本文主要关注“令牌选择”（token-choice）类型的MoE，因为它在生产环境中被广泛采用。我们使用以下符号：$L$表示一个步骤中被路由的令牌实例数量（例如，批量大小 × 序列长度），$K$表示每个令牌选择的专家数量，E表示专家总数，$d$表示模型维度。

2.1 门控网络与令牌路由

门控网络的功能。门控网络决定了每个输入令牌到最相关专家的路由。该网络通常是一个线性变换，将输入维度$d$映射到专家数量$E$，从而为每个令牌生成每个专家的得分。随后进行Top-K选择，为每个令牌收集门控得分最高的专家ID。门控输出定义为：topk experts = TopK(softmax(Wgx))，其中$W_g \in R^{E \times d}$是门控网络参数，$K$是每个令牌选择的专家数量。

传统路由方法的内存瓶颈。在Top-K选择之后，令牌必须被物理路由到相应专家的执行缓冲区。在传统实现中，此路由过程需要大量的辅助内存和额外的处理来压缩和存储分派的令牌，这构成了一个关键的内存瓶颈。

容量限制路由（Token-dropping）。早期的工作如Switch Transformers和GShard【【6】Fedus, W., Zoph, B., and Shazeer, N. Switch transformers: scaling to trillion parameter models with simple and efficient sparsity. J. Mach. Learn. Res., 2022.；【13】Lepikhin, D., Lee, H., Xu, Y., Chen, D., Firat, O., Huang, Y., Krikun, M., Shazeer, N., and Chen, Z. Gshard: Scaling giant models with conditional computation and automatic sharding. In ICLR 2021.】采用了容量限制的路由机制（令牌丢弃）。令牌按其门控得分排序并打包到专家e的缓冲区中；任何超过容量$C$的令牌要么被丢弃，要么被路由到残差路径。容量的典型选择是：

图片描述
图片描述

其中$\gamma$是用户定义的容量因子。容量限制路由因其固定大小的缓冲区而易于系统实现，但代价是模型质量下降。

无丢弃路由（Dropless Routing）。最近的文献【【16】Rajbhandari, S., Li, C., Yao, Z., Zhang, M., Aminabadi, R. Y., Awan, A. A., Rasley, J., and He, Y. DeepSpeedMoE: Advancing mixture-of-experts inference and training to power next-generation AI scale. In ICML, 2022.；【8】He, J., Qiu, J., Zeng, A., Yang, Z., Zhai, J., and Tang, J. Fastmoe: A fast mixture-of-expert training system. CoRR, 2021.】更关注无丢弃的路由机制，这通常能带来更好的模型质量。该方法确保每个令牌都由一个专家处理，从而提升模型质量并消除了容量因子调整的需求。然而，由于分配给每个专家的令牌数量是可变的，底层系统必须高效地管理动态的计算和内存需求。因此，近期文献主要集中在优化这些动态和可变长度工作负载的计算上【【7】Gale, T., Narayanan, D., Young, C., and Zaharia, M. Megablocks: Efficient sparse training with mixture-ofexperts. In Proceedings of Machine Learning and Systems, 2023.；【1】Aminabadi, R. Y., Holmes, C., Rajbhandari, S., Yao, Z., and He, Y. Turbomoe: Enhancing moe model training with smart kernel-fusion and data transformation. 2025.】。

持续存在的根本挑战。无论是令牌丢弃还是无丢弃路由方案，一个根本性的挑战依然存在：当前的实现需要存储索引和压缩后的令牌数据，导致内存占用与$L \times K \times d$成正比。在现代LLM训练中，随着序列长度和批量大小的增加，这导致内存缓冲区的急剧扩张。

内存占用量化示例。为了说明与令牌分派相关的激活内存占用，我们以一个真实的MoE模型（如DeepSeek）为例进行定量研究。对于DeepSeek模型中一个典型的MoE层，它有$L \approx 200$万个令牌，激活专家数$K = 4$，模型维度$d = 6144$，并为路由的令牌缓冲区使用每元素2字节（bfloat16），其内存占用为：

图1. 传统方法中的MoE与MoEBlaze的对比。左图展示了传统的MoE计算，包括令牌分发、专家计算和加权聚合（详见第2节）。右图展示了所提出的MoEBlaze框架，它引入了内存高效的令牌路由和专家计算（详见第3节）。
图1. 传统方法中的MoE与MoEBlaze的对比。左图展示了传统的MoE计算，包括令牌分发、专家计算和加权聚合（详见第2节）。右图展示了所提出的MoEBlaze框架，它引入了内存高效的令牌路由和专家计算（详见第3节）。
图片描述
图片描述

由此可见，仅一个路由缓冲区，单个MoE层就可能消耗近百GB的内存。

2.2 专家前馈网络（FFNs）

FFN计算过程。令牌分派之后是跨专家的前馈网络（FFN）计算。每个FFN计算通常实现为一个两层的多层感知机（MLP）。第一层将输入从维度$d$投影到更高维的隐藏空间$h$，第二层再投影回输出维度$d'$（假设$d'=d$）。所有$E$个专家的参数所需的总内存为$O(E \times (h \times d + d \times h))$，但条件计算范式确保每个令牌只有$k$个专家是激活的，从而保持了每次前向传播的低计算成本。专家$E_i$内对给定输入$x$的FFN计算定义为：

$$E_i(\mathbf{x})=W_{2,i} \cdot \sigma(W_{1,i}\mathbf{x})$$

其中$W_{1,i} \in R^{h \times d}$和$W_{2,i} \in R^{d \times h}$，$\sigma$是非线性激活函数（例如ReLU, GELU, SwiGLU）。

FFN中的内存瓶颈。第二个主要的内存瓶颈源于FFN计算过程中的中间激活存储。对于一个激活的专家，第一次线性变换$W_{1,i}x$会生成一个大小为$L_i \times h$的中间激活，其中$L_i$是路由到专家i的令牌数量。虽然激活的专家数量很少，但在前向传播期间，所有$E$个专家的这些激活的聚合内存为$O(L \times h)$，并且在反向传播期间可能更高，因为需要为梯度计算存储中间值。激活函数的选择（例如SwiGLU）将进一步加剧内存压力。

激活内存占用量化示例。我们使用deepseek的配置作为例子来说明FFN计算产生的激活内存占用的重要性。我们有$L \approx 200$万，FFN隐藏维度$h = 24576$，使用每元素2字节（bfloat16），我们可以得到中间激活的内存占用：

$$\text{Mem}_{act}=2L\times h\approx 98\text{GB}$$

2.3 输出聚合

聚合过程。最后阶段是输出聚合，其中所选专家的输出通过加权求和的方式组合，为每个令牌生成最终输出。权重来源于门控网络的分数。输入令牌$x$的MoE输出$y$为：

$$\mathbf{y} = \sum_{i=1}^{E} g_{i}(\mathbf{x}) \cdot E_{i}(\mathbf{x})$$

其中$g_i(x)$是专家$i$的门控分数，只有top-k个专家的分数非零。聚合输出所需的内存为$O(L \times d)$。计算上，这涉及$O(L \times k \times d)$次操作，考虑到稀疏性（$k \ll E$），这通常是高效的。

A2 方法细节

3 内存高效的令牌路由与训练算法

传统方法的缺陷。在令牌选择MoE中，门控网络将每个输入令牌分配给一个或多个专家。为了在专家计算期间方便地进行令牌索引和组织，传统系统将这些令牌压缩到每个专家的缓冲区中。这个压缩步骤，以及随后的每个专家MLP块的执行，会创建中间结果，这些结果在被求和并归约到原始令牌长度的输出之前，一直保持在压缩后的令牌长度。关键在于，这种分离和中间存储在整个MoE训练过程中引入了显著的激活缓冲区。

MoEBlaze的核心思想与流程。我们提出了一种内存高效的路由和专家计算算法，它能显著减少辅助激活占用，同时实现高效的MoE训练。算法的核心思想是利用在令牌分派步骤中生成的辅助索引列表，来跟踪路由决策，并在整个MoE计算过程中执行“即时”（on-the-fly）的令牌访问和结果归约。具体来说，我们的融合核（fused kernel）操作如下：
1. 它接收门控决策并构建专家-令牌索引列表及其他相关的索引结构。
2. 在专家-令牌索引列表的指导下，它通过从原始、未排列的激活张量中进行即时收集（gather）来执行专家MLP计算。
3. 然后，专家求和利用令牌-专家索引列表，直接将MLP结果求和并归约到最终的输出张量中。
通过直接访问输入并仅存储最终结果，我们消除了其他论文中通常需要用于物化令牌路由的许多中间激活。只存储令牌和专家ID的令牌-专家索引列表非常轻量。此外，这种方法使我们能够将令牌/专家索引与计算紧密融合，为重叠内存访问和计算提供了可能性。这在像最新的H100 GPU这样的现代硬件上特别有利，可以实现更好的资源利用和更快的速度。

3.1 前向传播

3.2 反向传播

反向传播的挑战与解决方案。反向传播接收($L, d$)令牌的梯度，并将其沿着前向步骤的逆向路径传播回去。传统的专家求和反向过程依赖于路由的令牌激活缓冲区，在通过MLP专家进行反向传播之前，执行一次“扩展”或物化，将($L, d$)的梯度扩展为($L \times k, d$)的“路由梯度令牌”。然而，我们提出的方法通过使用相同的反向映射索引来避免这个中间扩展步骤。

  1. 专家求和反向 (Expert Summation Backward)：利用从分派元数据派生的令牌映射结构，将($L, d$)梯度张量映射回($L \times k, d$)的路由梯度令牌。这是通过一个高效的操作完成的，该操作将输出梯度“散布”（scatter）到物化的中间MLP结果张量中的相应位置。
  2. 专家计算反向 (Expert Computation Backward)：接下来，梯度通过MLP向后流动。之前检查点化的两个背靠背MLP之间的中间结果将在这里用于计算权重梯度。
  3. 令牌梯度累积 (Token Gradient Accumulation)：最后，关于输入令牌的梯度从所有专家处累积起来。这一步对每个令牌路由到的k个专家的贡献进行求和，为输入激活生成最终的($L, d$)梯度张量。由于我们没有物化路由令牌结果的激活存储，我们也利用令牌索引数据结构来执行即时归约。

4 高效且可并行的分派与数据结构

4.1 数据结构

我们定义了上述内存高效MoE训练算法所需的关键数据结构。

图2. 用于高效内存MoE训练的数据结构。
图2. 用于高效内存MoE训练的数据结构。

示例。图2展示了一个$L=6$个令牌，$E=4$个专家，以及$k=2$个激活专家的示例的数据结构。从门控分数矩阵中，我们获得每个令牌的分配如下：
令牌0：专家{2, 3}；令牌1：专家{0, 1}；令牌2：专家{0, 3}；令牌3：专家{1, 2}；令牌4：专家{0, 3}。
将这些令牌的分配连接起来，我们将得到token_expert_indices为：

$\text{token_expert_indices} = [2, 3, 0, 1, 1, 2, 0, 3],$

同样，我们可以得到每个专家路由到的令牌：
专家0：令牌{1, 2, 4}；专家1：令牌{1, 3}；专家2：令牌{0, 3}；专家3：令牌{0, 2, 4}。
将它们连接起来，我们得到expert_token_indices和expert_token_offsets：

$$\begin{aligned} \begin{aligned} \text{expert\_token\_indices} &= [1, 2, 4, 1, 3, 0, 3, 0, 2, 4] \\ \text{expert\_token\_offsets} &= [0, 3, 5, 7, 10] \end{aligned} \end{aligned}$$

token_index_map存储了每个令牌在拼接后的专家令牌列表中的位置。例如，token_index_map[0] = {5, 7}，因为令牌0被路由到2个专家（$k=2$），并被放置在expert_token_indices的第5个和第7个位置。

4.2 高效的分派数据结构构建

构建过程的挑战。现在我们详细介绍高效构建上述数据结构的方法。构建过程面临一个挑战：以专家为中心的数据结构的内在设计需要一个多对一的映射，即多个令牌被分配给同一个专家。使用朴素的方法会导致在GPU架构上严重的线程级写冲突，从而影响性能。

基于排序的方法及其瓶颈。一种解决方案是依赖基于排序的方法来构建令牌分派。该方法将所有令牌的top-k选择结果（topk_experts）扁平化为一个长度为$Lk$的一维数组，其中包含（专家ID，令牌ID）元组。然后该数组按专家ID进行全局排序以对令牌进行分组，随后通过索引恢复来重构令牌顺序并计算每个专家的范围。这种排序过程虽然概念简单，但在规模较大时会引入严重的性能瓶颈。排序在GPU上实现为多趟基数排序（radix sort），需要多个与密钥宽度成比例的全局内存通道，迫使频繁的全局内存访问并多次移动$O(Lk)$的数据。这导致了实际上的高复杂度和GPU资源利用率低下。此外，这个全局排序步骤限制了细粒度的并行性，并强制采用具有高核函数启动延迟的多核函数分派流水线（多趟排序、分段扫描、索引恢复等）。

本文的高效构建方法。为此，我们引入了一种高效的方法，该方法用基于局部索引构建的可并行化方法替代了复杂的全局排序和组织，这种方法能很好地映射到GPU上。该方法是一个简单的三步过程，每一步都设计为无原子操作（atomic-free）且可在GPU上并行化，从而可以最小化昂贵的全局内存通道并避免复杂的多核函数流水线。

  1. 构建密集的令牌-专家映射 (Build Dense Token-Expert Map)：第一步，我们构建一个表示为dense_token_map的密集位图来编码top-k令牌到专家的路由。对于每个令牌i，我们考虑其top-k分配的专家{$e_{i,0}, \dots, e_{i,k-1}$}。对于每个门控槽，我们将dense_token_map[i, e_{i,k}]设置为i。所有其他条目保持未设置状态。这个编码映射的构建在GPU上是高度可并行的，我们为每个warp分配不相交的令牌行，保证了没有warp内冲突。
  2. 计算专家长度 (Compute Expert Lengths)：利用构建的dense_token_map，下一步是高效地计算每个专家的稀疏令牌ID列表的长度和偏移量。我们启动一个自定义核函数，其CTA网格映射到dense_token_map的列（专家）上。每个CTA专用于单个专家$e_i$，并计算该列中的非零条目（令牌到专家的分配）。使用warp级别的归约来聚合CTA内的行向总和，从而产生专家长度数组。expert_lengths[e_i]的值表示最终路由到专家$e_i$的令牌数量。长度计算之后，通过对expert_lengths数组应用前缀和来导出expert_offsets。
  3. 将索引路由到门控 (Route Indices to Gates)：第三步涉及生成每个专家的令牌ID列表expert_token_indices，它作为后续MLP计算的输入。为了在GPU上以无冲突的方式实现索引的紧凑、按专家拼接，我们采用了一个围绕生成位置图（location map）的两阶段过程。该图指定了dense_token_map中每个非零条目在expert_token_indices列表中的最终目标位置ID。一旦位置图构建完成，一个简单的并行核函数会从dense_token_map中读取元素，并将它们直接写入expert_token_indices中计算出的相应位置，从而保证了完全的并行性而无需原子操作。位置图的构建采用两步策略以确保其无原子操作构造：(i) 瓦片级扫描（tile-level scan）：我们为每个专家启动一个CTA。同一CTA内的线程处理分配给该专家的连续令牌。它们首先在共享内存中计算瓦片级别的计数，然后在CTA内部局部执行独占扫描操作（前缀和）。(ii) 得到的CTA局部独占扫描计数随后与专家预先计算的全局专家偏移量相加。这个加法操作得出了在拼接索引数组中的正确、最终的位置ID。

5 训练-核协同设计以实现端到端效率

本节详细介绍了我们联合优化混合专家（MoE）训练核和智能激活检查点方法，以解决与一些高级激活方法相关的内存问题。

5.1 SwiGLU MoE与内存瓶颈

高级激活函数的应用与挑战。现代MoE训练越来越多地采用先进的非线性激活函数，如SiLU和SwiGLU，以替代ReLU/GELU。先前的工作表明，这些激活函数提供了更平滑的非线性，可以改善优化稳定性，并在大规模语言任务上带来更好的经验准确性。虽然在数值上更有利，但这些激活函数在训练期间引入了更复杂的计算和更大的内存占用。

SwiGLU激活函数详解。我们以SwiGLU激活函数为例。SwiGLU激活定义为：

$$\text{SwiGLU}(\mathbf{x}; W_1, W_2) = \text{SiLU}(\mathbf{x}W_1) \cdot (\mathbf{x}W_2),$$

其中$SiLU(u) = u \cdot \sigma(u)$，$\sigma$是sigmoid操作。对于一个拥有$E$个专家的MoE层，每个专家实现一个SwiGLU前馈网络（FFN），一个路由到单个专家的令牌批次$x \in R^{L \times d}$会引发两个投影：$a = xW_1 \in R^{L \times h}$ 和 $b = xW_2 \in R^{L \times h}$。随后是逐元素操作$SiLU(a)$和最终的乘积$SiLU(a) \odot b$。

5.2 激活检查点与核协同设计

传统核的内存开销。在传统的核中，前向流水线需要物化多个中间结果以适应反向计算（例如，在SwiGLU示例中，这包括两个GEMM输出$a$和$b$，sigmoid $\sigma(a)$，$SiLU(a)$，以及最终的乘积）。这些中间结果被写入全局内存，随后又被读出，这会产生不可忽略的开销。随着模型和批量大小的扩展，这会产生显著的内存流量和存储成本，迅速成为一个不可忽视的瓶颈。

优化的核心观察。为了减轻观察到的内存压力，我们提出了一种对MoE训练流程及其底层GPU核的联合优化，该优化可以在不牺牲性能的情况下减少激活内存占用和内存流量。我们的优化基于以下观察：
* 在现代GPU上，激活函数的计算通常是内存带宽受限的。这主要有两个原因：1）激活函数的计算大多是逐点操作，而现代GPU对此类操作能力很强；2）在LLM训练中，我们通常处理的是令牌数量远大于嵌入维度（$L \gg d$）的情况。对这种“高瘦”形状的矩阵进行操作在GPU上通常是内存带宽受限的。
* 虽然激活计算在计算上很轻量，但其内存占用却惊人地大。对于复杂的现代激活函数尤其如此，它需要为反向传播物化和保存许多中间的、逐点的结果。由此产生的内存分配是巨大的，与批量大小、序列长度和FFN维度成线性关系。

联合激活检查点与核融合方法。基于这些观察，我们提出了联合激活检查点和核融合的方法。我们的方法将SwiGLU中的两个第一层投影与激活函数的后处理（epilogue）融合在一起，并在专门的路径中应用激活检查点，以“打破”由中间激活存储引起的“内存墙”。

后处理融合 (Epilogue Fusion)。为了减少激活流量和核启动开销，我们将两个第一层投影和SwiGLU后处理融合成一个单一的核。该核消耗非物化的路由令牌，只加载一次输入$x$，同时将其流式传输通过$(W_1, W_2)$两个GEMM，在寄存器/共享内存中计算$SiLU(a)$，并立即与$b$进行乘法运算，只将最终输出写入全局内存。这种“后处理融合”消除了$a, b$的全局写入以及随后的为逐元素操作的重新读取，有效地将计算从内存受限领域转移到尽可能大的计算受限领域。与为每个投影使用单独的核相比，它还将输入$x$的读取次数减半。

反向传播融合。在反向传播期间，融合两个第一层投影意味着必须聚合来自两条路径的关于共享输入$x$的梯度。我们的实现不是分配两个独立的激活缓冲区并将它们拼接起来，而是以融合的方式计算两个分支的激活导数，并通过瓦片式归约（tiled reductions）就地聚合梯度——完全消除了临时的全局缓冲区。

激活检查点策略。在融合的基础上，我们进一步应用了激活检查点策略——在前向传播期间我们跳过保存SwiGLU的中间结果（SiLU）。相反，我们在反向传播期间采用重新计算的策略，利用了SiLU函数计算成本低（例如逐元素操作）且在现代GPU上严重受内存带宽限制的事实。

算法 1 融合的SwiGLU MoE训练

输入: 输入令牌 X ∈ R^(L×d), 投影权重 W1, W2 ∈ R^(d×h), W3 ∈ R^(h×d)
输出: 输出 Y_out ∈ R^(L×d), 梯度 ∇W1, ∇W2, ∇W3, ∇Z

1: // Swiglu MoE训练的前向模块
2: Procedure FusedForward(X, W1, W2, W3)
3:   // 只加载一次输入令牌
4:   Load X
5:   // 第一个MLP投影:
6:   // 计算 A 和 B;
7:   // SiLU(A) 和 Y_swi 在核内计算
8:   // SiLU(A) 是瞬态的
9:   (A, B), Y_swi ← FusedSwiGLU(Z, W1, W2)
10:  Y_out ← Y_swi * W3
11:  Store A, B, Y_swi
12:  // 存储激活和SwiGLU输出以供反向传播使用
13:  Return Y_out
14:
15: Procedure FusedBackward(Y_out, ∇Y_out, W1, W2, W3, A, B, Y_swi)
16:  // 最终投影的梯度
17:  ∇W3 ← Y_swi^T * ∇Y_out
18:  // 将梯度反向传播到SwiGLU输出
19:  ∇Y_swi ← ∇Y_out * W3^T
20:  // 加载存储的激活
21:  Load A, B
22:  // 重新计算SiLU(A)以节省内存
23:  S_recomp ← SiLU(A)
24:  // 关于A的导数
25:  ∇A ← ∇Y_swi ⊙ B ⊙ ∇SiLU(A)
26:  // 关于B的导数
27:  ∇B ← ∇Y_swi ⊙ S_recomp
28:  ∇W1, ∇W2 ← FusedBwdW(X, ∇A, ∇B)
29:  ∇X ← FusedBwdX(∇A * W1^T, ∇B * W2^T)
30:  Return ∇W1, ∇W2, ∇W3, ∇Z

5.3 整合：端到端的Swiglu MoE训练

算法1总结了使用SwiGLU激活函数的MoE模型的端到端训练过程。伪代码具体展示了第5节中详述的激活检查点和核融合的集成。虽然融合核的底层实现为简洁起见被省略了，但其高层方法论源于第3节中解释的内存高效令牌分派。

A4 实验环境

我们在各种代表性的MoE配置中，将MoEBlaze与当前最先进的稀疏训练系统Megablocks【【7】Gale, T., Narayanan, D., Young, C., and Zaharia, M. Megablocks: Efficient sparse training with mixture-ofexperts. In Proceedings of Machine Learning and Systems, 2023.】进行基准测试，展示了在训练速度和内存效率方面的显著改进。

表1. 实验中使用的MoE配置。FFN隐藏维度设置为输入维度的四倍（ffn_hidden_size = 4 × input_d）。
表1. 实验中使用的MoE配置。FFN隐藏维度设置为输入维度的四倍（ffn_hidden_size = 4 × input_d）。

A4 实验结果

6.3 使用SiLU的MoE训练中的内存效率

如图3所示，在所有测试配置中，MoEBlaze相比Megablocks基线，持续且显著地减少了激活内存消耗。MoEBlaze实现的内存节省在具有大输入维度和高专家数量的配置中尤其显著，例如conf4。具体来说，MoEBlaze仅需6,100 MB内存，与Megablocks消耗的22,000 MB相比，实现了近3.6倍的减少。对于较小的配置（如conf1），激活内存节省不那么明显，这是预期的，因为节省量与序列长度L和激活专家数k成比例，而这两者在conf1中都很小（k=1）。峰值激活内存的大幅减少是两个核心系统创新的直接结果：（1）一种更高效的令牌分派机制，最小化了中间缓冲区的分配；（2）在我们的自定义激活检查点方案中采用了智能重计算。

图3. 在表1定义的MoE配置集上，使用SiLU激活函数时MoEBlaze和Megablocks的激活内存占用比较。
图3. 在表1定义的MoE配置集上，使用SiLU激活函数时MoEBlaze和Megablocks的激活内存占用比较。

6.4 基于SiLU的MoE训练速度

图4展示了在给定配置下，MoEBlaze相对于Megablocks的训练速度提升。MoEBlaze取得了显著的性能增益，加速比在1.4倍到3.7倍之间。最大加速在conf4（$D_{input}$ = 2048, E = 16, L = 1024, B = 32）上实现，表明MoEBlaze在较大模型维度下扩展性特别好。这些训练速度的提升归因于三个因素：（1）我们高度优化的令牌分派实现，减少了与昂贵的令牌分派和置换操作相关的延迟开销；（2）高效的数据分派构建核，它非常轻量级且在GPU上运行迅速，避免了其他基于排序方法中昂贵的多趟核，并极大地消除了CPU端的瓶颈；（3）用于批处理GEMM计算的融合核，有效利用了H100最新的硬件加速特性，如warp-group矩阵乘法、张量内存加速器等。

图4. 在表1的配置集上，使用SiLU作为激活函数时，MoEBlaze相对于Megablocks的加速比。
图4. 在表1的配置集上，使用SiLU作为激活函数时，MoEBlaze相对于Megablocks的加速比。

6.5 使用SwiGLU的MoE训练中的内存效率

SwiGLU激活函数由于额外的门控和逐元素乘法操作，天生需要更高的内存使用量。图5显示了在SwiGLU设置下的内存消耗比较。MoEBlaze仍然保持了对Megablocks的巨大内存优势，其峰值激活内存通常不到基线使用量的一半。例如，在配置conf3中，Megablocks需要超过40,000 MB，而MoEBlaze则控制在约10,000 MB。这种持续4倍的内存压力降低证实了我们的内存高效分派和智能重计算方案即使对于更复杂的激活函数也同样高效。

图5. 在表1定义的MoE配置集上，使用SwiGLU激活函数时MoEBlaze和Megablocks的激活内存占用比较。
图5. 在表1定义的MoE配置集上，使用SwiGLU激活函数时MoEBlaze和Megablocks的激活内存占用比较。

6.6 基于SwiGLU的MoE训练速度

图6展示了使用SwiGLU激活时MoEBlaze相对于Megablocks的加速比。与ReLU的结果相比，加速因子通常更高且更一致，范围从2倍到6.2倍。相对速度的增加是两个因素的结果：（1）SwiGLU中更复杂的计算为MoEBlaze高度融合的核提供了更多超越基线的机会；（2）我们的激活优化带来的内存带宽节省在SwiGLU情况下更为关键，因为其中间激活尺寸更大、更复合，从而通过智能核融合和重计算减少了过多的全局内存访问，使得MoEBlaze能够更快地执行整个核。

图6. 在表1的配置集上，使用SwiGLU作为激活函数时，MoEBlaze相对于Megablocks的加速比。
图6. 在表1的配置集上，使用SwiGLU作为激活函数时，MoEBlaze相对于Megablocks的加速比。

A7 补充细节 (相关工作)

A5 结论与未来工作

我们提出了MoEBlaze，一个在GPU上进行MoE训练的快速且内存高效的系统。MoEBlaze通过融合的令牌分派和计算核设计，消除了物化大型每专家激活缓冲区的需求。此外，MoEBlaze整合了MoE计算和激活流水线，以最小化读/写流量，从而更好地节省内存带宽和减少内存占用。我们的实验表明，MoEBlaze在一系列配置中提供了高效且可扩展的解决方案，峰值激活内存消耗减少了4倍以上，端到端训练加速达到了6.2倍。

虽然本文主要关注单设备性能，但我们注意到MoEBlaze的核心机制也适用于分布式设置。作为未来的工作，我们计划将MoEBlaze扩展到分布式训练框架，并研究多节点、多GPU的MoE训练优化。