GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding

文章标题：GShard: 通过条件计算和自动分片扩展巨型模型
作者/机构：Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, Zhifeng Chen (Google)

A1 主要贡献

本文旨在解决扩展神经网络时面临的计算成本、编程易用性和并行设备实现效率等挑战。

• 核心问题：随着模型规模的扩大，现有的深度学习框架在模型并行化方面支持不足，计算成本呈超线性增长，巨型模型的计算图表示和编译成为瓶瓶颈，且实现高效的分区策略需要大量非凡的工程努力。

• 研究目标：设计一个系统，能够高效地训练参数量远超单个加速器内存容量的巨型模型（例如，数千亿参数），同时保持训练成本的亚线性增长和良好的编程易用性。

• 主要贡献/创新点：

  1. GShard模块：本文提出了GShard，一个由轻量级标注API和XLA编译器扩展组成的模块。它允许模型开发者通过简单的标注来表达复杂的并行计算模式，而无需修改大量现有模型代码，从而将模型描述与并行化实现解耦。
  2. 条件计算与模型架构：将稀疏门控混合专家（Sparsely-Gated Mixture-of-Experts, MoE）层与Transformer架构相结合，设计了一种能够以亚线性计算成本扩展模型容量的架构。通过这种方式，成功构建并训练了一个超过6000亿参数的多语言神经机器翻译模型。
  3. 可扩展的SPMD编译器：开发了一种基于单程序多数据（SPMD）的编译器技术，该技术为所有设备生成一个统一的程序，使得编译时间与设备数量无关，解决了传统多程序多数据（MPMD）方法在扩展到数千设备时面临的编译瓶颈。
  4. 大规模应用验证：在2048个TPU v3加速器上，仅用4天时间就成功训练了6000亿参数的巨型模型。该模型在100种语言到英语的翻译任务上取得了远超现有技术的质量。如图1所示，模型从375亿参数扩展到6000亿参数（16倍），翻译质量持续提升，而端到端训练成本仅从6个TPU v3核心年增加到22个（3.6倍），实现了成本的亚线性增长。

A3 设计原则

设计巨型模型训练框架的思路。本文通过构建一个6000亿参数的序列到序列Transformer模型，并结合稀疏门控混合专家层（Sparsely-Gated Mixture-of-Experts layers），展示了如何克服扩展模型的挑战。该模型实现了亚线性的计算成本增长和$O(1)$的编译时间。我们在2048个TPU v3设备上训练此模型4天，在100种语言到英语的翻译任务上，使用单一非集成模型取得了远超现有技术的翻译质量。实验发现，随着模型变大，翻译质量提高，但总训练时间相对于模型大小呈亚线性增长，如图1所示。为了构建如此巨大的模型，我们做出了以下关键设计选择。

亚线性扩展。首先，模型架构的设计应确保计算和通信需求在模型容量方面呈亚线性增长。条件计算【25, Conditional computation in neural networks for faster models, 2015; 16, Outrageously large neural networks: The sparsely-gated mixture-of-experts layer, 2017; 26, Depth-adaptive transformer, 2020; 27, Controlling computation versus quality for neural sequence models, 2020】通过在每个输入的基础上激活一个子网络，使我们能够满足训练和推理的效率要求。通过添加Position-wise稀疏门控混合专家（MoE）层【16, Outrageously large neural networks: The sparsely-gated mixture-of-experts layer, 2017】来扩展基于RNN的机器翻译和语言模型的容量，已经以亚线性的计算成本取得了最先进的结果。因此，我们在第2节中介绍了将Transformer架构扩展到MoE层的方法。

抽象的力量。其次，模型描述应与分区实现和优化分离。这种关注点分离让模型开发者能够专注于网络架构并灵活更改分区策略，而底层系统则应用保留语义的转换并实现高效的并行执行。为此，我们提出了一个模块GShard，它仅要求用户在模型中用分区策略标注几个关键张量。它由一组简单的标注API和XLA【28, XLA: Optimizing Compiler for TensorFlow, 2019】中的编译器扩展组成。模型开发者编写模型时，就好像只有一个拥有巨大内存和计算能力的设备，而编译器会根据标注和自身的启发式方法自动为目标设备分区计算。我们在第3.2节中提供了更多标注示例。

可扩展的编译器。第三，系统基础设施，包括计算表示和编译，必须能够扩展到数千个设备以进行并行执行。例如，图2说明了在4个设备上（用颜色编码）划分点积操作的两种不同方式。请注意，使用通常的MPMD（多程序多数据）方法（图2a），随着设备数量的增加，图中的节点数线性增加，扩展变得更具挑战性。相反，我们开发了一种用于SPMD（单程序多数据）转换的编译器技术，该技术生成一个在所有设备上运行的单一程序，使编译时间与设备数量无关，保持恒定，如图2b所示。我们将在第3.3节中更详细地讨论我们的SPMD框架。

A2 方法细节

2. 模型架构

2.1 Transformer架构的稀疏扩展

MoE层替换FFN层。Transformer架构【10, Attention is all you need, 2017】已广泛用于自然语言处理，并成为许多序列到序列任务（如机器翻译）的事实标准。Transformer利用编码器和解码器两个计算块，两者都通过堆叠多个Transformer层实现。Transformer编码器层由两个连续的层组成，即一个自注意力层后跟一个位置前馈层。解码器增加了第三个交叉注意力层，该层关注编码器的输出。我们通过条件计算对Transformer进行稀疏扩展，方法是在编码器和解码器中，将每隔一个前馈层替换为带有top-2门控变体的Position-wise混合专家（MoE）层【16, Outrageously large neural networks: The sparsely-gated mixture-of-experts layer, 2017】（图3）。我们通过改变Transformer层的数量和每个MoE层中专家的数量来扩展模型容量。

子网络激活与计算成本的亚线性扩展。每个训练样本由一对子词标记序列组成。在训练和推理过程中，每个标记都会激活MoE Transformer的一个子网络。该子网络的大小大致与每个MoE层中的专家数量无关，从而实现了前一节所述的计算成本的亚线性扩展。计算复杂性在第3.1节中进一步分析，训练性能在第5节中分析。

2.2 Position-wise混合专家层

MoE层定义。我们模型中使用的混合专家（MoE）层基于【16, Outrageously large neural networks: The sparsely-gated mixture-of-experts layer, 2017】的设计，但在稀疏门控函数和辅助损失方面有所变化。用于Transformer的MoE层由E个前馈网络$FFN_1, ..., FFN_E$组成：

门控函数的目标。门控函数GATE(·)对MoE层至关重要，它通过一个softmax激活函数来建模，以指示每个专家在处理输入标记时的权重，换言之，指示一个专家处理输入标记的能力。此外，门控函数必须满足两个目标：
- 负载均衡：希望MoE层能为给定标记稀疏地激活专家。一个简单的解决方案是根据softmax概率分布选择top-k个专家。然而，已知这种方法会导致训练中的负载不平衡问题【16, Outrageously large neural networks: The sparsely-gated mixture-of-experts layer, 2017】：训练中看到的大多数标记会被分派到少数专家，导致少数（繁忙的）专家积累了非常大的输入缓冲区，而其他专家未得到训练，从而减慢了训练速度。同时，许多其他专家根本没有得到充分的训练。一个更好的门控函数设计应能更均匀地在所有专家之间分配处理负担。
- 规模化效率：如果门控函数是顺序执行的，实现负载均衡将相当简单。对于输入批次中的所有N个标记和E个专家，仅门控函数的计算成本至少为$O(N \cdot E)$。然而，在我们的研究中，N的数量级为百万，E的数量级为千，顺序实现的门控函数将使大部分计算资源在大部分时间内处于空闲状态。因此，我们需要一个高效的并行门控函数实现来利用众多设备。

门控函数的实现机制。我们设计了GATE(·)函数中的以下机制来满足上述要求（详见算法1）：
- 专家容量（Expert capacity）：为确保负载均衡，我们强制规定一个专家处理的标记数量低于某个统一的阈值，我们将其定义为专家容量。假设训练批次中的总标记数为N，每个标记最多分派给两个专家，则专家容量设为$O(N/E)$。GATE(·)维护一个运行计数器$c_e$，记录分派给一个专家的标记数量。当一个标记选择的两个专家都已超出其容量时，该标记被视为溢出标记，此时$G_{s,E}$退化为零向量。这些标记的表示$x_s$通过残差连接传递到下一层。
- 本地分组分派（Local group dispatching）：GATE(·)将训练批次中的所有标记平均分成G个组，即每组包含$S=N/G$个标记。所有组都独立并行处理。每组被分配每个专家的部分容量，即$2N/(G \cdot E)$。每组确保最多只有这么多标记被分派给一个专家。通过这种方式，我们可以确保专家容量仍然得到执行，并且整体负载是均衡的。
- 辅助损失（Auxiliary loss）：重要的是门控函数不能总是选择相同的少数专家，因为这会导致少数专家的容量溢出和其余专家的利用率不足。遵循【16, Outrageously large neural networks: The sparsely-gated mixture-of-experts layer, 2017】，我们定义了一个辅助损失项$\aux$来强制执行此约束。它以一个常数乘数k加到模型的总损失函数中：$L = \nll + k * \aux$。算法1第(13)行中辅助损失项$\aux$的具体形式是基于以下考虑：项$c_e/S$表示路由到每个专家的输入分数，我们希望最小化$c_e/S$的均方。但由于$c_e$是从top-2操作派生的，不可微，我们使用每个专家的平均门控值$m_e$作为可微的近似，并将$(c_e/S)^2$替换为$m_e(c_e/S)$，这样就可以通过梯度下降进行优化。
- 随机路由（Random routing）：直观地说，因为$y_s$是所选专家返回值的加权平均，如果第二个专家的权重非常小，我们可以简单地忽略第二个专家以节省整体专家容量。因此，除了遵守专家容量约束外，GATE(·)还以与其权重$g_2$成比例的概率分派给第二优的专家。

3. 使用GShard实现高度并行化

本节描述了第2节中模型的实现，该模型能在TPU设备集群上高效运行。

实现步骤概述。第一步是将模型表示为线性代数操作，我们的软件栈（TensorFlow【21, Tensorflow: a system for large-scale machine learning, 2016】）和硬件平台（TPU）对此进行了高度定制和优化。将模型的大部分内容以线性代数形式编码是相当容易的，与原始Transformer的方式相同。然而，由于其顺序性，表达MoE层，特别是算法1中呈现的GATE(·)函数，需要一些努力，我们在3.1节中描述了细节。

通过标注表达并行性。接下来，我们对线性代数计算进行标注以表达并行性。计算中的每个张量都可以使用3.2节中的分片API进行标注，以在设备集群中复制或分布。使用分片标注可以实现模型描述和高效并行实现之间的关注点分离，并允许用户灵活地表达各种并行化策略。例如，(1) 注意力层通过沿批次维度拆分并将其权重复制到所有设备来进行并行化。另一方面，(2) MoE层中的专家由于其庞大的规模，无法在所有设备上复制，唯一可行的策略是将专家分片到许多设备上。此外，整个模型在这两种模式(1)-(2)之间交替。使用标注将模型开发者从系统优化工作中解放出来，并避免将并行实现和底层细节固化到模型代码中。

编译器自动并行化。最后，编译器基础设施接收一个（部分）标注的线性代数计算，并生成一个可扩展到数千个设备的高效并行程序。如3.3节所述，编译器应用SPMD（单程序多数据）分区变换来表达每个设备的计算，插入必要的跨设备通信，处理不规则模式（如不均匀分区），并最终生成一个在所有设备上启动以进行并行执行的单一程序。

3.1 Positions-wise混合专家层的线性代数表示

用张量运算表示MoE层。我们的模型实现（算法2）将整个加速器集群视为单个设备，并将其核心数学算法表示为几个与集群具体设置无关的张量操作。爱因斯坦求和表示法【30, Die grundlage der allgemeinen relativitätstheorie, 1923】（即tf.einsum）是简洁表达模型的强大结构，我们在实现中广泛使用了它。softmax门控计算可以通过一个einsum后跟softmax函数轻松表示。将输入分派给选定专家的操作通过分派掩码和输入之间的单个einsum表示。所有$FFN_e$的权重被组合成单个3-D张量$w_i$和$w_o$，而$FFN_1, ..., FFN_E$的计算则使用3个操作符（两个einsum和一个relu）表示。最后，将所有专家输出的加权平均值合并到最终输出中，由另一个einsum表示。

Top2Gating实现细节。算法2中的Top2Gating计算了算法1中描述的所有组局部$G_{S,E}$的并集。combine_weights是一个形状为[G, S, E, C]的4-D张量。当g组中的输入标记s被发送到专家e的输入缓冲区的c位置时，combine_weights[g, s, e, c]的值非零。对于特定的g和s，切片combine_weight[g, s, :, :]最多包含两个非零值。二进制的dispatch_mask是通过将combine_weights中所有非零值设为1来生成的。

算法可扩展性分析。我们需要适当地选择组数G和专家数E，以便算法可以扩展到具有D个设备的集群。分析在给定一个包含N个标记的训练批次下，一个训练步骤的总体计算复杂性（浮点运算总数）是值得的。

Algorithm 2: Forward pass of the Positions-wise MoE layer. The underscored letter (e.g., G and E ) indicates the dimension along which a tensor will be partitioned.

1 gates = softmax ( einsum ("GSM ,ME - >GSE", inputs , wg ))   
2 combine_weights , dispatch_mask = Top2Gating ( gates )   
3 dispatched_expert_inputs = einsum (   
4 "GSEC ,GSM - >EGCM ", dispatch_mask , reshaped_inputs )   
5 h = einsum ("EGCM ,EMH - >EGCH ", dispatched_expert_inputs , wi )   
6 h = relu ( h )   
7 expert_outputs = einsum ("EGCH ,EHM - >GECM ", h , wo )   
8 outputs = einsum (   
9 "GSEC ,GECM - >GSM", combine_weights , expert_outputs )

计算复杂度分析。我们基于以下假设分析算法2的计算复杂度随设备数D的扩展情况：a) 每个设备的标记数$N_D = O(1)$是常数；b) $G = O(D)$, $S = O(1)$ 且 $N = O(GS) = O(D)$；c) $M = O(1)$, $H = O(1)$；d) $E = O(D)$；e) $C = O(2S/E) = O(1/D)$，其中$D < S$且为正整数。算法2中的总浮点运算数(FLOPS)为：

通信成本分析。除了计算成本，我们还有非恒定的跨设备通信成本，但当我们增加D时，它以一个适度的速率$O(\sqrt{D})$增长（第5节）。

3.2 GShard并行执行的标注API

分片API概述。由于算法1中张量的巨大规模和计算需求，我们必须在许多设备上并行化该算法。算法2中带下划线的字母说明了如何对每个张量进行分片。GShard中的分片API允许我们标注程序中的张量，以选择性地指定它们应如何分区。此信息会传播到编译器，以便编译器可以自动应用转换以进行并行执行。我们在工作中使用了TensorFlow/Lingvo【31, Lingvo: a modular and scalable framework for sequence-to-sequence modeling, 2019】中的以下API。
- replicate(tensor)：标注张量在所有分区上复制，并返回标注后的张量。这通常用于我们模型中的非MoE层以复制权重。
- split(tensor, split_dimension, num_partitions)：标注张量沿split_dimension进行分区，并返回标注后的张量。分区i被放置在第i个设备上，num_partitions不能超过系统上的设备数。
- shard(tensor, device_assignment)：泛化了split()，允许对多个维度进行分区并指定每个分区的位置。附录A.3更详细地描述了这个API。
请注意，对split或shard的调用仅添加标注，不会改变用户程序中的逻辑形状。用户仍然使用完整的形状，无需担心不均匀分区等问题。

GShard的通用性。GShard是通用的，因为简单的API以相同的方式应用于所有维度。根据用例的不同，分片的维度可以包括批次（数据并行）、特征、专家，甚至图像模型中的空间维度。此外，由于分片标注是针对每个张量的，模型的不同部分可以以不同的方式进行分区。这种灵活性使我们能够对巨大的MoE权重进行分区，并在MoE和非MoE层之间切换分区模式，以及支持本文之外的用例，例如大图像的空间分区【32, Train ML models on large images and 3D volumes with spatial partitioning on Cloud TPUs, 2019】（附录A.4）。

代码示例。通过上述分片API，我们可以将算法2中显示的分片策略表示如下。输入张量沿第一个维度拆分，门控权重张量被复制。在计算分派的专家输入后，我们应用split将分片从组（G）维度更改为专家（E）维度。D是设备数量。

# Partition inputs along group (G) dim .
+ inputs = split ( inputs , 0 , D )
# Replicate the gating weights
+ wg = replicate ( wg )
gates = softmax ( einsum ("GSM ,ME - >GSE", inputs , wg )) 
6 combine_weights , dispatch_mask = Top2Gating ( gating_logits )
dispatched_expert_inputs = einsum ( 
8 "GSEC ,GSM -> EGCM ", dispatch_mask , reshaped_inputs ) 
9 # Partition dispatched inputs along expert (E) dim. 
10 + dispatched_expert_inputs = split ( dispatched_expert_inputs , 0 , D ) 
11 h = einsum ("EGCM ,EMH - > EGCH ", dispatched_expert_inputs , wi )

逐张量分片分配。如上例所示，用户无需对程序中的每个张量进行标注。标注通常只需在几个重要的操作符上，如我们模型中的Einsum，编译器会使用自己的启发式方法推断其余张量的分片。例如，由于输入张量沿G分区而权重张量被复制，编译器选择沿相同的G维度对einsum输出进行分区（第5行）。类似地，由于两个输入都沿G维度分区用于输入分派einsum（第7行），输出分片被推断为沿G维度拆分，然后我们在输出上添加split标注以重新分片到E维度。上述示例中的一些标注也可以由编译器确定（例如，replicate(wg)），但建议对计算的初始输入和最终输出张量进行标注。

编译器推断逻辑。编译器目前使用迭代数据流分析，从用户标注的操作符开始，将分片信息从一个操作符传播到其邻居（操作数和用户）。该分析试图通过对齐相邻操作符的分片决策来最小化重新分片的机会。可能还有其他方法，如整数规划或机器学习方法，但改进自动分片分配不是本文的重点，我们将其作为未来工作。

混合手动与自动分片。带有分片标注的自动分区通常足以应对常见情况，但GShard也具有灵活性，允许将手动分区的操作符与自动分区的操作符混合使用。这为用户提供了更多关于如何分区操作符的控制，一个例子是用户拥有超出操作符语义的更多运行时知识。例如，XLA或TensorFlow的Gather操作符定义都没有传达关于输入中不同范围的索引边界信息，但用户可能知道某个特定的Gather操作符仅在每个分区内洗牌数据。在这种情况下，用户可以通过简单地缩小维度大小并执行本地Gather来轻松地对操作符进行分区；否则，编译器将需要对索引范围持保守态度并增加不必要的通信开销。例如，算法2中第3行的分派Einsum，它使用一个one-hot矩阵来分派输入，可以替代地用一个Gather操作符实现，使用简单的手动分区，而模型的其余部分则自动分区。以下是说明此用例的伪代码。

1 # input has shape [G, S, M]. split () does not change logical shape .   
2 input = split ( input , 0 , num_devices )   
3 # s_indices has shape [E, G, C, 1]. Values : indices to S in input .   
4 s_indices = split ( s_indices , 1 , num_devices )   
5   
6 # Begin manual partitioning .   
7 # partitioned_input has shape [G/ num_devices , S, M]   
8 partitioned_input = auto_to_manual_spmd_partition ( input )   
9 # partitioned_s_indices has shape [E, G/ num_devices , C, 1]   
10 partitioned_s_indices = auto_to_manual_spmd_partition ( s_indices )   
11 # Concat with G indices in partitioned_input : Iota on G dimension .   
12 partitioned_gs_indices = concat (   
13 iota ([ E , G / num_devices , C , 1] , 1) , partitioned_s_indices , 3)   
14 # partitioned_data has shape [E, G/ num_devices , C, M]   
15 partitioned_data = gather (   
16 partitioned_input , partitioned_gs_indices )   
17   
18 # Switch back to auto partitioning .   
19 # data has shape [E, G, C, M]   
20 data = manual_to_auto_spmd_partition ( partitioned_data )   
21 ...

3.3 GShard的XLA SPMD分区器

分区器概述。本节描述了根据分片标注自动对计算图进行分区的编译器基础设施。分片标注告知编译器每个张量应如何在设备间分布。SPMD（单程序多数据）分区器（或简称“分区器”）是一个编译器组件，它将一个计算图转换为一个在所有设备上并行执行的单一程序。这使得编译时间几乎与分区数量无关，从而使我们能够扩展到数千个分区。

在XLA中实现。我们在XLA编译器【28, XLA: Optimizing Compiler for TensorFlow, 2019】中实现了该分区器。包括TensorFlow、JAX、PyTorch和Julia在内的多个前端框架已经有将它们的图表示降低到XLA HLO图的逻辑。与TensorFlow等流行的前端框架相比，XLA的操作符集要小得多，这减轻了实现分区器的负担，同时不损害通用性，因为现有的前端降低逻辑完成了使其具有表达力的繁重工作。尽管我们在XLA中开发了该基础设施，但我们在此描述的技术可以应用于其他机器学习框架中的中间表示（例如，ONNX【33, ONNX: Open Neural Network Exchange, 2019】、TVM Relay【34, Relay: a new ir for machine learning frameworks, 2018】、Glow IR【35, Glow: Graph lowering compiler techniques for neural networks, 2018】）。

分区器核心。XLA将计算建模为一个数据流图，其中节点是操作符，边是操作符之间流动的张量。分区器的核心是逐操作处理，根据输入和输出上指定的分片，将一个全尺寸操作符转换为一个分区尺寸的操作符。当计算被分区时，会引入各种跨设备数据传输的模式。为了在大型规模下最大化性能，定义一组核心通信原语并为目标平台优化它们至关重要。

3.3.1 通信原语

通信模式。由于分区器强制所有设备运行相同的程序，通信模式也是规则的，XLA定义了一组执行MPI风格通信【36, MPI: A Message-Passing Interface Standard, 2009】的集体操作符。下面列出了我们在SPMD分区器中使用的常见通信原语。
- CollectivePermute：该操作符指定一个源-目标对列表，源的输入数据被发送到相应的目标。它用于两个地方：改变分片张量在分区间的设备顺序，以及本节后面讨论的halo交换。
- AllGather：该操作符按照指定的顺序连接所有参与者的张量。它用于将分片张量更改为复制张量。
- AllReduce：该操作符对所有参与者的输入执行元素级归约（例如，求和）。它用于组合来自不同分区的部分归约的中间张量。在TPU设备网络中，当分区数量增加时，AllReduce具有恒定的成本（第5.2节）。它也是一种常用的原语，在其他类型的网络拓扑中也有高效的实现【37, BlueConnect: Decomposing All-Reduce for Deep Learning on Heterogeneous Network Hierarchy, 2019】。
- AllToAll：该操作符逻辑上将每个参与者的输入沿一个维度拆分，然后将每个片段发送给不同的参与者。在从其他参与者接收到数据片段后，每个参与者连接这些片段以产生其结果。它用于将分片张量从一个维度重新分片到另一个维度。AllToAll是TPU设备网络中进行这种重新分片的有效方式，其成本随分区数量的增加而亚线性增长（第5.2节）。

3.3.2 逐操作符的SPMD分区

分区器核心转换。分区器的核心是根据指定的分片，将一个全尺寸操作符逐个转换为分区尺寸操作符。虽然某些操作符（如元素级操作）的支持是微不足道的，但我们讨论了需要跨分区通信的几种常见情况。

技术挑战。在一般情况下存在一些重要的技术挑战，我们将在第3.3.3节中介绍。为了使讨论更贴近MoE模型，本节重点介绍Einsum分区，以说明几种通信模式。为了简化，我们暂时假设所有张量都是均匀分区的，即要分区的维度大小是分区数的倍数。

Einsum案例研究。Einsum是实现MoE模型中最关键的操作符。它们在XLA HLO中表示为Dot操作，其中每个操作数（LHS或RHS）包含三种类型的维度：
- 批次维度：这些是易于并行的维度。LHS、RHS和输出中必须存在相同的批次维度集，并且输出中的每个元素仅依赖于LHS和RHS中相应的批次。
- 收缩维度：仅存在于操作数中。LHS和RHS必须具有相同的收缩维度集，它们在输出中被求和并折叠。
- 非收缩维度：也是并行维度，存在于其中一个操作数和输出中。LHS和RHS各自有自己的非收缩维度集，这些维度被输出继承。

Einsum分区策略。分片传播优先选择在LHS、RHS和输出的批次维度上使用相同的分片，因为这可以避免任何跨分区通信。然而，这并非总是可能的，我们在以下三种情况下需要跨分区通信。
- 重分片（Resharding）：在我们构建的MoE模型中，专家分派逻辑（算法2中的第3行）需要在Einsum之后切换分区维度。由于使用AllToAll进行重分片是高效的（第5.2节），我们首先在本地执行Einsum，然后将其重分片到期望的维度，如图4a所示。
- 累积部分结果（Accumulating partial results）：如果输入是沿着收缩维度分区的，本地结果是部分的，我们需要使用AllReduce来组合它们并产生最终结果，如图4b所示。
- 循环中切片（Slicing in a loop）：对于某些场景，我们还实现了一种类似于Cannon算法【38, A Cellular Computer to Implement the Kalman Filter Algorithm, 1969】的算法，以限制每个分区上的张量大小。例如，如果两个操作数都在非收缩维度上分区，我们无法直接计算本地Einsum，因为操作数具有不同的非收缩维度。复制其中一个操作数不会导致冗余计算，但它要求被复制的操作数能够放入设备内存。因此，如果操作数的大小太大，我们转而保持两个操作数都分区，并使用一个循环来迭代结果的每个切片，并使用CollectivePermute来通信输入切片（图4c）。

3.3.3 支持完整的操作符集

通用性挑战与解决方案。我们解决了一些额外的挑战，以使SPMD分区器能够支持完整的操作符集，而无需对张量形状或操作符配置施加额外约束。这些挑战通常涉及分区之间的非对称计算或通信模式，这在SPMD中尤其难以表达，因为单个程序需要对所有分区都具有通用性。我们不能简单地基于运行时设备ID在单个程序中创建许多分支，因为这会导致程序大小的爆炸。

静态形状与不均匀分区。XLA要求张量形状是静态的。然而，当一个计算被分区时，并非所有分区都具有相同的输入/输出形状，因为维度可能不能被分区数整除。在这些情况下，形状的大小向上取整到分区数的下一个倍数，并且该填充区域中的数据可以是任意的。当计算一个操作符时，我们可能需要为正确性在填充区域填充一个已知值。例如，如果我们需要对一个Reduce-Add操作符进行分区，需要使用零作为单位元。考虑一个例子，其中分区维度（15）不能被2（分区数）整除，因此分区1比需要的列多一列。我们创建一个范围为[0, 8)的Iota操作符，加上分区偏移量（从PartitionId × 8计算），并与完整形状偏移量（15）进行比较。基于谓词值，我们从操作数或零中选择，结果是被掩码的操作数。

静态操作符配置。XLA操作符具有静态配置，如Convolution中定义的填充、步长和膨胀。然而，不同的分区可能不会以相同的操作符配置执行。例如，对于一个Convolution，最左边的分区在其左侧应用填充，而最右边的分区在其右侧应用填充。在这种情况下，分区器可能会选择一些配置，使得某些分区产生比需要稍多的数据，然后切掉不相关的部分。附录A.4讨论了Convolution和类似操作符的例子。

Halo交换。某些操作符具有一种通信模式，涉及与相邻分区进行部分数据交换，我们称之为halo交换。我们使用CollectivePermute操作符在分区之间交换halo数据。halo交换最典型的用例是用于分区基于窗口的操作符（例如，Convolution、ReduceWindow），因为相邻分区可能需要重叠的输入数据（图5a）。实际上，由于窗口配置（膨胀、步长和填充）的高级使用以及不均匀的halo大小，这些操作符的halo交换通常需要与适当的填充、切片和掩码相结合。我们在附录A.4中描述了各种场景。halo交换的另一个用途是用于改变形状大小的数据格式化操作符。例如，在Slice或Pad操作符之后，张量的形状发生变化，分区之间的边界也随之改变。这要求我们重新对齐不同分区上的数据，这可以作为一种halo交换形式来处理（图5b）。其他数据格式化操作符，虽然逻辑上不改变形状的大小，但也可能需要halo交换，特别是因为静态形状约束和不均匀分区。例如，Reverse操作符反转张量中元素的顺序，但如果它被不均匀地分区，我们需要在分区之间移动数据以保持填充在逻辑上位于结果张量的右侧。另一个例子是Reshape。考虑将一个[3, 2]的张量重塑为[6]，其中输入在第一个维度上以2种方式不均匀分区（分区形状[2, 2]），输出也以2种方式分区（分区形状[3]）。由于不均匀分区，输入上有填充，但在Reshape之后，输出张量不再有填充；因此，需要以类似于Slice的方式进行halo交换（图5c）。

编译器优化。SPMD分区器创建了各种数据格式化操作符以执行切片、填充、连接、掩码和halo交换。为了解决这个问题，我们利用XLA在TPU上的融合能力，以及对切片和填充的代码移动优化，来在很大程度上隐藏数据格式化的开销。因此，即使对于大量使用掩码和填充的卷积网络，运行时开销通常也可以忽略不计。

A4 实验

实验环境

• 数据集：

  • 名称：内部网络规模数据集。
  • 规模与用途：从网络爬取的并行文档，覆盖100种语言与英语的互译，总计250亿训练样本。实验主要关注100种语言到英语的翻译，约使用130亿训练样本。数据具有噪声大、领域多样、语言间样本量极不均衡（从数十亿到数万不等）的特点。

• 模型架构：

  • 模型：稀疏门控混合专家Transformer（MoE Transformer）。
  • 关键参数：
    • 模型维度：1024
    • 前馈网络（FFN）和MoE隐藏层维度：8192
    • 多头注意力头数：16
    • 注意力键/值维度：128
    • Dropout率：0.1
  • 配置：测试了不同深度（L=12, 36, 60层）和专家数量（E=128, 512, 2048）的组合，参数量从300亿到6000亿不等（见表1）。还测试了一个1万亿参数的模型，但因数值稳定性问题未报告结果。
  • 基线模型：1) 为100个语言对分别训练的双语NMT模型；2) 一个使用GPipe训练的96层密集Transformer模型（T(96L)）。

• 硬件配置：

  • 加速器：最多2048个TPU v3核心。
  • 连接：TPU设备通过2D环形网络连接。

• 软件配置：

  • 框架：TensorFlow / Lingvo【31, Lingvo: a modular and scalable framework for sequence-to-sequence modeling, 2019】。
  • 编译器：XLA。
  • 优化器：Adafactor【88, Adafactor: Adaptive learning rates with sublinear memory cost, 2018】，学习率为1.0，并在1万步后采用平方根衰减。
  • 数据类型：模型权重和激活使用float32以保证训练稳定性。
  • 分词器：SentencePiece【89, Sentencepiece: A simple and language independent subword tokenizer and detokenizer for neural text processing, 2018】，源语言使用包含102种语言的64k词表，目标语言（英语）使用32k词表。

实验结果

翻译质量 (M4)

实验旨在通过大规模多语言、大规模模型（M4）来验证GShard的有效性。通过将模型扩展到处理100种语言，目标是同时提升高资源和低资源语言的翻译质量。

• 模型扩展效果显著：如图6所示，所有MoE Transformer模型在翻译质量（∆BLEU）上均显著优于各语言对的独立双语基线。随着模型规模的增加，质量持续提升。最大的6000亿参数模型（MoE(2048E, 36L)）取得了最佳的平均翻译质量。

• 深度带来普遍增益：在专家数量固定的情况下，增加模型深度（例如，从12层增加到36层）能为所有语言（无论高资源还是低资源）带来一致且显著的质量提升，平均提升2-3个BLEU点（表3）。这表明更深的模型具有更强的泛化能力。

• 专家数量缓解容量瓶颈：增加专家数量（即增加模型宽度和参数量）能显著缓解高资源语言的容量瓶颈问题。例如，将专家数从128增加到512，平均BLEU提升了3.3。然而，从512增加到2048时，增益下降到1.3，显示出收益递减的趋势。这表明容量瓶颈主要存在于128到512个专家之间。

• 密集模型在低资源语言上表现更佳：与稀疏的MoE模型相比，96层的密集Transformer模型（T(96L)）在低资源语言上表现出更强的正向迁移能力（图6中右侧）。这可能是因为密集模型的所有参数都是共享的，为低资源任务提供了更大的知识迁移带宽。

训练效率

• 样本效率：更深的模型样本效率更高。如表2所示，深度为3倍的MoE模型（例如36L vs 12L）达到相同的训练损失所需的训练令牌数减少了2到3倍。这验证了过参数化对加速收敛的积极作用。

• 训练时间：GShard的训练效率极高。最大的6000亿参数模型在2048个TPU v3核心上仅需4天即可完成训练（处理1T tokens），总成本为22.4个TPU核心年。相比之下，性能稍逊的96层密集模型需要超过10倍的训练时间（235个TPU核心年），而训练100个双语基线模型也需要29个TPU核心年（表3）。这证明了条件计算在扩展模型时的成本优势。

性能与内存消耗

• 内存效率：GShard实现了优秀的内存可扩展性。如图7所示，当增加设备和专家数量时，每个设备的内存消耗大致保持恒定（$O(1)$）。对于层数非常多的模型，编译器会自动启用重计算（rematerialization）来减少峰值激活内存，例如MoE(2048E, 60L)的激活内存小于MoE(2048E, 36L)。

• 运行时效率：训练步时（step time）随模型规模呈亚线性增长。当模型规模从128个专家扩展16倍到2048个专家时，执行时间仅增加了1.7倍。如图8所示，模型实现了较高的硬件利用率（roofline performance），在128专家时达到>70%，在2048专家时仍能达到48%。

• 通信开销：通信是影响扩展性的关键。微基准测试（图9）表明，AllReduce操作在TPU网络上的开销与设备数无关（$O(1)$），而用于重分片的AllToAll操作开销呈亚线性增长（约$O(\sqrt{D})$）。这些高效的通信原语是GShard能够高效扩展MoE模型的关键。

A5 结论

本文介绍了GShard，一个通过轻量级分片标注实现大规模计算自动分区的深度学习模块。GShard提供了易于使用且灵活的API，用于扩展巨型神经网络。我们应用GShard将带有稀疏门控专家（MoE）层的Transformer架构扩展，并展示了一个6000亿参数的多语言神经机器翻译模型，它可以在4天内高效训练，在100种语言到英语的翻译任务中，用单一模型取得了优于现有技术的性能和质量。

主要发现与经验：
1. 持续扩展的有效性：结果表明，持续扩展神经网络能带来一致的质量提升，证明模型质量尚未随着规模扩大而达到瓶颈。
2. 训练效率的重要性：条件计算等算法改进，结合易用的接口，可以有效地利用大规模计算能力，在规模和计算成本之间取得有利的权衡。这使得训练万亿参数级别模型的实验周期从数月或数周缩短到数天。
3. 抽象与可扩展编译器的价值：将模型描述与并行化实现分离的抽象层，以及生成单一通用程序的SPMD分片编译器，是实现可扩展编译的关键。
4. 对模型容量的再思考：实验结果支持，单纯的参数数量并不总能与模型的有效容量相关联，尤其是在大规模多任务学习和数据不平衡的场景下。

总结而言，本文的研究阐明了一条现实可行的神经网络扩展路径，即模型扩展应与训练效率齐头并进，通过算法、系统和硬件的协同设计，实现模型质量的持续提升。

A6 附录

A.1 使用扁平化束搜索进行解码

解码过程。在解码时，我们使用类似于【61, Google’s neural machine translation system: Bridging the gap between human and machine translation, 2016】的带长度归一化的束搜索（beam search）。解码是自回归的，一次生成一个目标序列标记，因此对于长度为m的输出，解码器层堆栈会顺序执行m次。特别是对于每个解码器MoE层，都有分派/组合操作，这需要跨设备通信。推理使用与训练相同的集群和相同数量的设备。

扁平化束搜索实现。在束搜索过程中，我们将束假设（beam hypotheses）扁平化为一个包含所有底层标记交错排列的单一序列，并修改解码器自注意力掩码，以便每个假设只关注联合扁平序列中的适当位置。我们对每个解码器自注意力层维护的键/值张量应用相同的变换。这使我们能够避免在每次束扩展后重新排序先前计算的注意力键/值。相反，我们只重新排序表示当前活动假设的0/1掩码。然而，注意力的长度会变为k倍。

性能权衡。这种权衡可能是积极的也可能是消极的，取决于实现细节。正如【87, Fast transformer decoding: One write-head is all you need, 2019】中所解释的，内存带宽限制对于使用Transformer模型进行增量解码非常重要。从这个角度看，通过扁平化束，我们用一个计算/内存比较高的单一操作（在更长序列上进行注意力点积，键更多）替换了两个计算/内存比较低的操作（注意力点积和键/值重排序），但它必须访问的总内存量是相同的。

A.3 通用分片API

shard() API。除了第3.2节中列出的两个常用分片API（replicate()和split()）之外，用户或编译器可以使用更高级的分片策略来最小化数据传输。shard(tensor, device_assignment)用提供的设备分配来标注张量进行分区，并返回标注后的张量。我们使用设备分配（一个多维整数数组）来表示分区的完成方式。device_assignment的秩与数据张量相同；其元素数量是分区的总数，每个元素是占据相应数据片的设备的ID。例如，一个形状为[3, 16, 64]的3D张量，设备分配形状为[1, 2, 4]，将具有分区形状[3, 8, 16]，并且设备分配中元素的顺序决定了每个分区占据哪个片。

考虑设备拓扑。由于跨设备的数据移动严重影响并行执行性能，因此在为device_assignment中的设备ID分配时，考虑目标设备拓扑以及张量分区之间的通信以获得最佳性能非常重要。图10显示了基于设备拓扑和张量上行式通信模式的两种不同设备分配。

A.4 卷积和基于窗口的操作符的SPMD分区

空间维度分区。GShard能够对卷积中的空间维度进行分区，并且足够通用以支持像巨型图像这样的用例【32, Train ML models on large images and 3D volumes with spatial partitioning on Cloud TPUs, 2019】。要对卷积层进行空间分片，我们可以按以下方式使用分片API。

GShard随后会将空间维度上的分片传播到其他层和反向传播过程。本节的其余部分讨论了对卷积和类似操作符进行分区的具体复杂性。有几种基于窗口的操作（例如，卷积、ReduceWindow），它们都需要某种类型的halo交换，因为数据可能在窗口之间共享。我们使用CollectivePermute操作符在分区之间交换halo数据，但一个复杂之处在于不同分区的halo大小可能不同，而CollectivePermute需要是静态形状的。

窗口配置。我们首先介绍SPMD分区器必须考虑的窗口配置。卷积中的每个空间维度都有以下配置集。
- Stride（步长）：窗口移动以产生下一个输出元素的距离（元素数量）。
- Low/high padding（低/高填充）：在LHS（基底）维度的低/高端填充的元素数量。
- Base dilation（基底膨胀）：LHS的膨胀因子，即每个元素之间填充的元素数量加一（不包括低/高填充）。无基底膨胀意味着值为1。
- Window dilation（窗口膨胀）：RHS（窗口）中每个元素之间填充的元素数量加一。

非恒定halo大小。我们用一个没有膨胀的简单例子来证明非恒定halo大小是常见的。图11显示了一个4路分区的卷积，其中分区的右halo大小为(1, 2, 3, 4)，可以表示为分区ID的线性函数：partition_id + 1。分区1负责生成2个输出元素（红色单元格），这意味着该分区需要从分区0获取0个元素，从分区2获取2个元素（由两个虚线红色窗口覆盖的区域）。

通用halo交换序列。图12描述了通用halo交换的操作序列。首先，我们计算所有分区中左、右halo的最大尺寸，并执行最大尺寸的halo交换（步骤1和2）。由于某些分区可能有多余的halo，我们使用DynamicSlice（基于分区ID）来切掉当前分区的有效区域（步骤3）。最后，某些分区可能包含垃圾值（例如，来自范围外输入数据的halo），因此我们应用第3.3.3节中描述的掩码。

基底膨胀处理。基底膨胀为halo交换增加了额外的复杂性，因为每个分区的偏移量可能位于膨胀孔中，而且低/高填充是在膨胀后应用的，这使得边缘的行为与内部元素不同。我们分3种情况处理基底膨D胀（图13）。
- 情况1：stride × per_shard_window_count可被dilation整除，其中per_shard_window_count是每个分区要处理的窗口数。此条件保证所有分区在LHS的第一个数据元素之前都以相同数量的（内部或低）填充元素开始，因此我们可以使用相同的低填充。
- 情况2：stride == 1但per_shard_window_count不能被dilation整除。在这种情况下，不同分区上的低填充是不同的。幸运的是，当stride == 1时，填充和膨胀后的基底区域上的所有位置都是有效的窗口起点，我们可以对所有分区使用最大的低填充，以确保每个分区计算所有必需的窗口，然后在分区窗口操作的输出上执行DynamicSlice以删除不必要的数据。
- 情况3：stride != 1且stride × per_shard_window_count不能被dilation整除。如果以上两个条件都不满足，不同分区可能以不同数量的填充元素开始，并且并非所有偏移量都是有效的窗口起点。解决这个问题的方法是在填充基底区域之外再填充窗口。

窗口膨胀处理。如果RHS被复制，窗口膨胀仅在基于其LHS对操作符进行分区时影响有效窗口大小。如果膨胀后的RHS也被分区（通常发生在跨步卷积的梯度计算中），处理窗口膨D胀仍然比处理基底膨胀简单，因为RHS上没有低/高填充。我们省略了实现的细节。