分片矩阵及其乘法

分区记法与集合通信操作

当我们在一万块 TPU 或 GPU 上训练 LLM 时，从抽象层面看，所做的计算与在一块设备上训练时仍然相同。区别在于，我们的数组无法装入单块 TPU/GPU 的 HBM，因此必须把它们拆开。¹ 我们把这种操作称为对数组进行“分片”或“分区”。实现扩展的艺术，就在于弄清楚如何对模型分片，同时保持计算高效。

下面是一个跨 4 块 TPU 分片的二维数组 A：

请注意，分片数组仍然拥有与未分片数组相同的全局形状或逻辑形状，例如 (4, 128)；但它还具有设备局部形状，例如 (2, 64)，这个形状给出了每块 TPU 实际保存的数据字节数（在上图中，每块 TPU 保存整个数组的 ¼）。接下来，我们把这一概念推广到任意数组。

统一的分片记法

我们使用一种命名轴记法的变体，来描述张量如何以数据块形式分片到各台设备上：假设存在一个称为设备网格的二维或三维设备栅格，并且每条轴都具有网格轴名称，例如 X、Y 和 Z。 然后，我们可以通过描述数组的每个命名维度如何沿物理网格轴分区，来指定矩阵数据在设备网格上的布局。我们把这种指派称为一种分片方案。

示例（上图）：对于上图，我们有：

• 网格： 上面的设备网格 Mesh(devices=((0, 1), (2, 3)), axis_names=('X', 'Y'))，它告诉我们共有 4 块 TPU，排列为 2x2 网格，两条轴分别命名为 $X$ 和 $Y$。
• 分片： $A[I_X, J_Y]$，它告诉我们将第一条轴 $I$ 沿网格轴 $X$ 分片，并将第二条轴 $J$ 沿网格轴 $Y$ 分片。该分片方案说明，每个分片保存数组的 $1 / (\lvert X\rvert \cdot \lvert Y\rvert)$。

综合这两点，我们知道该数组的局部形状（单台设备所保存分片的大小）为 $(\lvert I\rvert / 2, \lvert J\rvert / 2)$，其中 $\lvert I\rvert$ 是 A 第一维的大小，$\lvert J\rvert$ 是 A 第二维的大小。

小测验［沿 1 条轴进行二维分片］： 考虑一个数组 fp32[1024, 4096]，其分片为 $A[I_{XY}, J]$，网格为 {'X': 8, 'Y': 2}。每台设备保存多少数据？在 H100 上从 HBM 加载这个数组需要多长时间（假设每块芯片的内存带宽为 3.4e12）？

将这些分片方案可视化： 我们来观察一个被拆分到 4 台设备上的二维数据数组，尝试直观理解这些分片方案：

我们把矩阵的完全复制形式直接写成 $A[I, J]$，不附加任何分片指派。这意味着每台设备都包含整个矩阵的一份完整副本。

我们可以用下标形式的网格轴来表示其中某个维度已沿一条网格轴分区。例如，$A[I_X, J]$ 表示逻辑轴 I 已沿网格维度 X 分区，但维度 J 没有分区，因此这些数据块仍在网格轴 Y 上部分复制。

$A[I_X, J_Y]$ 表示逻辑轴 I 已沿网格轴 X 分区，而维度 J 已沿网格轴 Y 分区。

下图展示了其他可能性：

这里，$A[I_{XY}, J]$ 表示我们把网格轴 X 和 Y 看作一个更大的展平维度，并将命名轴 I 分区到所有设备上。多个网格轴下标的顺序很重要，因为它指定了在网格上遍历分区的顺序。

最后请注意，我们不能让多个命名轴沿同一个网格维度分片。例如，$A[I_X, J_X]$ 是一种没有意义且被禁止的分片。一旦某个网格维度被用于分片数组的一个维度，它在某种意义上就已经“用掉”了。

小测验： 设 A 是一个形状为 int8[128, 2048]、分片为 $A[I_{XY}, J]$、网格为 Mesh({'X': 2, 'Y': 8, 'Z': 2})（总计 32 台设备）的数组。每台设备上的 A 使用多少内存？跨所有设备计算，A 总共使用多少内存？

如何用代码描述这些分片？

到目前为止，我们一直避而不谈代码，但现在正好可以先睹为快。JAX 使用的命名分片语法与我们上面描述的抽象语法非常接近。我们会在第 10 节详细讨论这一点，不过这里先做一个快速预览。你可以在 Google Colab 中运行这里的代码，并对结果进行性能分析，看看 JAX 如何处理不同的分片方案。这段代码完成 3 件事：

1. 创建一个 jax.Mesh，把 8 块 TPU 映射为 4x2 网格，并将名称 ‘X’ 和 ‘Y’ 分别指派给两条轴。
2. 创建矩阵 A 和 B，其中 A 的两个维度都被分片，B 则沿输出维度分片。
3. 编译并执行一个简单的矩阵乘法，返回分片数组。

import jax
import jax.numpy as jnp
 
# Create our mesh! We're running on a TPU v2-8 4x2 slice with names 'X' and 'Y'.
# The Auto axis type tells JAX to let the XLA compiler infer intermediate shardings.
assert len(jax.devices()) == 8
Auto = jax.sharding.AxisType.Auto
mesh = jax.make_mesh(axis_shapes=(4, 2), axis_names=('X', 'Y'), axis_types=(Auto, Auto))
 
# A little utility function to help define our sharding. A PartitionSpec is our
# sharding (a mapping from axes to names).
def P(*args):
  return jax.NamedSharding(mesh, jax.sharding.PartitionSpec(*args))
 
# We shard both A and B over the non-contracting dimension and A over the contracting dim.
A = jnp.zeros((8, 2048), dtype=jnp.bfloat16, device=P('X', 'Y'))
B = jnp.zeros((2048, 8192), dtype=jnp.bfloat16, device=P(None, 'Y'))
 
# We can perform a matmul on these sharded arrays! out_shardings tells us how we want
# the output to be sharded. JAX/XLA handles the rest of the sharding for us.
y = jax.jit(lambda A, B: jnp.einsum('BD,DF->BF', A, B), out_shardings=P('X', 'Y'))(A, B)

import jax
import jax.numpy as jnp
 
# Create our mesh! We're running on a TPU v2-8 4x2 slice with names 'X' and 'Y'.
# The Auto axis type tells JAX to let the XLA compiler infer intermediate shardings.
assert len(jax.devices()) == 8
Auto = jax.sharding.AxisType.Auto
mesh = jax.make_mesh(axis_shapes=(4, 2), axis_names=('X', 'Y'), axis_types=(Auto, Auto))
 
# A little utility function to help define our sharding. A PartitionSpec is our
# sharding (a mapping from axes to names).
def P(*args):
  return jax.NamedSharding(mesh, jax.sharding.PartitionSpec(*args))
 
# We shard both A and B over the non-contracting dimension and A over the contracting dim.
A = jnp.zeros((8, 2048), dtype=jnp.bfloat16, device=P('X', 'Y'))
B = jnp.zeros((2048, 8192), dtype=jnp.bfloat16, device=P(None, 'Y'))
 
# We can perform a matmul on these sharded arrays! out_shardings tells us how we want
# the output to be sharded. JAX/XLA handles the rest of the sharding for us.
y = jax.jit(lambda A, B: jnp.einsum('BD,DF->BF', A, B), out_shardings=P('X', 'Y'))(A, B)

JAX 很酷的一点是，这些数组的行为就像它们根本没有分片一样！B.shape 会告诉我们全局形状或逻辑形状 (2048, 8192)。我们必须真正查看 B.addressable_shards，才能知道它在局部如何分片。我们可以对这些数组执行操作，JAX 会尝试判断该如何广播或重塑它们以完成操作。例如在上面的示例中，A 的局部形状是 [2, 1024]，B 的局部形状是 [2048, 4096]。JAX/XLA 会根据需要，自动在这些数组之间加入通信，以完成最终的乘法。

使用分片数组进行计算

如果有一个分布在许多设备上的数据数组，并希望对它执行数学运算，那么将数据与计算同时分片会带来哪些额外开销？

显然，这取决于具体计算。

• 对于逐元素运算，在分布式数组上执行操作没有额外开销。
• 当我们希望跨多个设备上的元素执行运算时，情况就复杂了。好在对于大多数机器学习任务，几乎所有计算都以矩阵乘法的形式进行，而矩阵乘法相对容易分析。

本节其余部分将讨论如何对分片矩阵做乘法。近似来看，这需要搬动矩阵的数据块，以便完整地对每个数据块执行乘法或求和。每种分片方案都会涉及不同的通信。 例如，$A[I_X, J] \cdot B[J, K_Y] \to C[I_X, K_Y]$ 不需要任何通信即可相乘，因为收缩维度（J，也就是我们实际求和的维度）没有分片。但是，如果希望输出不分片（即 $A[I_X, J] \cdot B[J, K_Y] \to C[I, K]$），就需要把 $A$ 和 $B$，或者把 $C$，复制到每台设备上（使用 AllGather（全聚合））。这两种选择的通信成本不同，因此需要计算成本并选择较低者。

方便的是，所有可能的分片方式大致可以归纳为 4 种需要考虑的情形；对于每种情形，需要加入何种通信都有一条相应规则

1. 情形 1： 两个输入都未沿收缩维度分片。我们可以直接对局部分片相乘，不需要任何通信。
2. 情形 2： 一个输入的收缩维度已分片。通常，我们会沿收缩维度对这个分片输入执行“AllGather”。
3. 情形 3： 两个输入都沿收缩维度分片。我们可以先对局部分片相乘，然后对结果执行“AllReduce”。
4. 情形 4： 两个输入各有一个非收缩维度沿同一条轴分片。若不先对其中一个输入执行 AllGather，我们就无法继续。

你可以把它们当成必须遵守的规则，但理解这些规则为何成立、成本多高也同样很有价值。下面我们将逐一详细讨论。

情形 1：两个乘数的收缩维度均未分片

引理： 在分片矩阵相乘时，只要收缩维度没有分片，并且两个矩阵没有沿同一条轴分片，那么计算就是有效的，输出也会沿用输入的分片方案。例如，下面的计算完全没有问题

它完全不需要通信，所得张量沿 X 和 Y 两个硬件维度分片。请试着想想为什么。基本原因是，计算与分片是相互独立的：每个批次条目都拥有收缩轴的某个局部数据块，可以在这个数据块上完成乘法与归约。以下任一种情形都能正常工作，并遵循这条规则：

由于 A 和 B 的收缩维度 J 都没有分片，我们只需执行输入的局部分块矩阵乘法，结果就会已经按照期望的输出分片方案完成分片。当两个乘数的非收缩维度沿同一条轴分片时，情况便不再如此（详见无效分片一节）。

情形 2：一个乘数的收缩维度已分片

考虑输入 A 沿收缩维度 J 分片，而 B 完全复制时该怎么办：

我们不能直接把 A 和 B 的局部数据块相乘，因为需要对 A 的完整收缩维度求和，而这个维度已沿 X 轴拆分。通常，我们先对 A 的分片执行“AllGather”，让每台设备都获得完整副本，然后才与 B 相乘：

这样，实际乘法就可以完全在每台设备上完成。

要点： 当相乘的矩阵中有一个沿收缩维度分片时，我们通常先对它执行 AllGather，使收缩维度不再分片，然后执行局部矩阵乘法。

请注意，当 B 也没有沿 X 分片时，我们还可以执行局部的部分矩阵乘法，然后对分片的部分和求和（也就是执行 AllReduce（全归约））；这样可以把计算分片，但通常通信成本更高。在某些情况下，这种方法可能更快，不过实践中 B 通常也会分片。下方的问题 4会推导它何时更优。

什么是 AllGather？ AllGather 是我们要讨论的第一个核心 MPI 通信原语。AllGather 会移除沿某条轴的分片，把散布在各台设备上的分片重新组装到该轴上的每台设备中。使用上面的记法，AllGather 会从一组轴中移除下标，例如

我们不必移除某个维度的全部下标。例如，$A[I_{XY}, J] \rightarrow A[I_Y, J]$ 也是 AllGather，只是仅沿一条轴执行。还请注意，我们也可能希望使用 AllGather 移除非收缩维度上的分片，例如在下面的矩阵乘法中：

我们既可以一开始就对 A 执行 AllGather 来移除输入分片，也可以先执行分片矩阵乘法，再对结果 C 执行 AllGather。

AllGather 实际是如何执行的？ 要围绕单条 TPU 轴（一个环）执行一维 AllGather，我们基本上会让每块 TPU 沿环传递自己的分片，直到每台设备都拥有一份副本。² 下面是一段动画：

我们可以单向执行 AllGather，也可以双向执行（上图展示的是双向）。单向执行时，每块 TPU 都会沿环传送大小为 $\text{bytes} / N$ 的数据块，共经过 $N - 1$ 跳。双向执行时，则要经过 $\lfloor \frac{N}{2} \rfloor$ 跳，每跳的数据大小为 $2 \cdot \text{bytes} / N$。

这需要多长时间？ 我们以双向 AllGather 为例，计算它的耗时。设 $V$ 为数组的字节数，$X$ 为收缩维度上的分片数。根据上图，每一跳都会在每个方向上发送 $V / \lvert X\rvert$ 字节，因此每跳耗时为

其中 $W_\text{ici}$ 是双向 ICI 带宽。³ 为了抵达每块 TPU，总共需要发送 $\lvert X\rvert / 2$ 跳⁴，因此整个归约过程耗时为

请注意，这个时间与 $X$ 无关！ 这多少有些令人惊讶，因为它意味着，尽管 TPU 只进行局部连接，连接的局部性却并不重要。我们的瓶颈只是每条链路的速度。

要点： 在吞吐量受限的场景中执行 AllGather（或者 ReduceScatter 或 AllReduce）时，实际通信时间只取决于数组大小和可用带宽，而与数组分片所跨的设备数量无关！

关于 ICI 延迟的一点说明： 无论数据量有多大，每次跨越 ICI 链路的一跳都有某种固有开销，通常约为 1us。这意味着，当数组 $A$ 很小、每跳耗时低于 1us 时，就会进入“延迟受限”状态，此时计算时间确实取决于 $X$。

下面是在 TPU v5e 8x16 切片上测得的 AllGather 带宽。数组沿大小为 16 的轴分片，因此拥有一个完整的双向环。

请注意，我们不仅达到了标称峰值带宽（4.5e10）的约 95%，而且在大约 10MB 时就达到了这个峰值；经过 16 路分片后，每台设备约为 625kB（顺带一提：这比 GPU 好得多）。

沿多条轴执行 AllGather 时会怎样？ 沿多条轴聚合时，我们可以在 ICI 的多个维度上执行聚合。例如，AllGather_XY([B, D_XY]) 会跨两条硬件网格轴操作。这会使可用带宽提高 $N_\text{axes}$ 倍。

考虑延迟时，我们会得到下面这条通用规则：

其中 $\sum_i \lvert X_i \rvert / 2$ 是 TPU 网格中最长路径的长度。

小测验 2［AllGather 时间］： 使用第 2 部分中的数值，在网格为 {'X': 8, 'Y': 4} 的 TPU v5e 上，对 bfloat16 格式、$E = 2048$、$F = 8192$ 的数据执行 AllGather_Y([E_Y, F]) → [E, F] 需要多长时间？若 $E=256, F=256$，又需要多长时间？

注意： 当我们拥有 {'X': 16, 'Y': 4} 这样的二维网格时，并不要求每条轴都对应某条特定的硬件轴。例如，上述网格可以描述一个 4x4x4 的 TPU v5p 立方体，其中 $X$ 轴映射了 2 条硬件轴。稍后讨论跨多条轴的数据并行时，这一点会派上用场。

情形 3：两个乘数的收缩维度均已分片

第三种基本情形是，两个乘数的收缩维度都沿同一条网格轴分片：

在这种情况下，局部的分片分块矩阵乘法至少是可以执行的，因为它们共享同一组收缩索引。但是，每个乘积只代表完整目标乘积的一个部分和，并且沿 X 维度的每台设备都会留下这个最终目标乘积的不同部分和。这种情况极其常见，因此我们扩展记法，显式标记这种状态：

记法 { U_X } 读作“沿网格轴 X 尚未归约”，表示该运算在某种意义上仍“不完整”，还要等待最后一次求和才能完成。$\cdot_\text{LOCAL}$ 语法表示我们执行局部求和，但让结果保持未归约状态。

这可以看作矩阵乘法与外积之间的如下关系：

其中 ⊗ 表示外积。因此，如果轴 X 上的 TPU i 保存 A 的第 i 列和 B 的第 i 行，我们就可以执行局部矩阵乘法，得到 $A_{:,i} \otimes B_{i,:} \in \mathbb{R}_{n\times m}$。在这个矩阵的每个条目中，都保存着 A • B 对应条目求和式的第 i 项。我们仍需对 P 执行这次求和——而 P 已沿网格轴 X 分片——才能得到完整的 A • B。如果把 A 和 B 写成分块形式（即分片），再对所得结果的每个分片求和，原理也完全相同。

我们可以沿 X 轴执行一次完整的 AllReduce 来完成这次求和：

AllReduce 会消除部分和，使沿该轴的每台设备都得到相同的、已完全求和的值。AllReduce 是本节将讨论的几种关键通信操作中的第二种；第一种是 AllGather，另外两种是 ReduceScatter 和 AllToAll。AllReduce 接收一个包含未归约（部分求和）轴的数组，让这些分片沿未归约轴传递并累加结果，从而完成求和。其签名为

这意味着，它只是移除 $\\{U_Y\\}$ 后缀，结果的其他部分保持不变。

AllReduce 的成本有多高？ 可以这样理解 AllReduce 的执行过程：每台设备把自己的分片发送给邻居，并把收到的所有分片相加。显然，这比 AllGather 更昂贵，因为每个“分片”的形状都与完整数组相同。一般而言，AllReduce 的成本是 AllGather 的两倍。 一种理解方式是，AllReduce 可以表示为另外两种原语的组合：一次 ReduceScatter（归约分散） 加一次 AllGather。与 AllReduce 一样，ReduceScatter 会消除数组中的部分和，但会产生沿给定维度“分散”或分区的输出。AllGather 会收集所有这些片段，并沿该物理轴“取消分区/取消分片/复制”相应逻辑轴。

那么 ReduceScatter 呢？ 正如 AllGather 会重新组装分片数组（移除一个下标），ReduceScatter 会对尚未归约/部分求和的数组求和，然后沿同一条网格轴分散（分片）另一条逻辑轴。$X[F]\\{U_Y\\} \to X[F_Y]$。动画展示了具体过程：请注意，它与 AllGather 十分相似，只是我们不再保留每个分片，而是把它们相加。因此，除去执行归约本身所需的时间，其延迟大致相同。

每一跳的通信时间，就是每分片字节数 $V / Y$ 除以带宽 $W_\text{ici}$，与 AllGather 相同，因此有

其中 $W_\text{ici}$ 是双向带宽，前提是我们拥有一个可用于归约的完整环。

情形 4：两个乘数各有一个非收缩维度沿同一轴分片

在对张量分片时，每个网格维度最多只能出现一次。应用上述规则有时会产生违反这条约束的情况，例如：

这种分片无效，因为沿维度 X 的某个分片（比如 i）将保存 C 的第 (i, i) 个分片，也就是一个对角条目。这样一来，所有分片合在一起也没有足够信息来恢复除对角条目以外的任何内容，因此我们不能允许这种分片。

解决办法是对其中某些维度执行 AllGather。这里有两种选择：

或者

无论哪种情况，结果的形状都只会提到一次 X。具体选择哪一种，取决于后续运算需要什么分片方案。

深入理解 TPU 通信原语

前面的 4 种情形引出了执行分片矩阵乘法时会用到的几种“核心通信原语”：

1. AllGather： 从分片中移除一个下标，收集各个分片。
2. ReduceScatter： 通过对某条轴上的分片求和，从数组中移除“尚未归约”后缀，同时让数组沿另一条轴保持分片状态。
3. AllReduce： 移除“尚未归约”后缀，使数组沿该轴不再分片。

还有一种核心通信原语值得一提，它会出现在专家混合（Mixture of Experts，MoE）模型及其他计算中：AllToAll（全互换）。

最后一种通信原语：AllToAll

最后一种基础集合通信在研究分片矩阵乘法时并不会自然出现，但在实践中却始终会遇到，它就是 AllToAll；更准确地说，是分片转置或重新分片操作这一特例。例如

在分片计算的不同区域采用彼此不兼容的布局方案时，通常需要使用 AllToAll 来重新排列分片布局。它会自然地出现在分片专家混合模型中。你可以把 AllToAll 理解为把一个下标从一条轴移到另一条轴。因为 AllToAll 不需要在环上复制每个分片的全部数据，所以它实际上比 AllGather 更便宜（成本为后者的 ¼）。⁵

推广到 ND AllToAll 后，对于一个跨所有设备合计有 $V$ 字节、位于 AxBxC 网格上的数组，总成本为

其中，照例 $W_\text{ici}$ 表示双向 ICI 带宽，而 $N = A \cdot B \cdot C \cdot \ldots$ 是设备总数。等价地，使用每台设备的字节数 $V / N$ 表示，成本为 $(V / N) \cdot \max(A, B, C, ...) / (4 \cdot W_\text{ici})$。对于一维网格，它化简为 $V / (4 \cdot W_\text{ici})$，即 AllGather 成本的 1/4。在二维情况下，成本实际上会随较短轴的大小增加而下降。

题外话：如果你想粗略推导这一结论，可以从一维环面 $\mathbb{Z} / N\mathbb{Z}$ 开始。随机选取源节点和目标节点时，两者的平均距离为 N / 4 跳，成本因此为 $(V \cdot N) / (4 * N)$。再考虑 ND 环面时，各条轴基本彼此独立。每个节点拥有 $1 / N$ 字节，平均需要让自己的数据跳转 $\max(A, B, C, …) / 4$ 跳。也可以从对分带宽推导：在 AllToAll 中，网格的每一半都向另一半发送自己一半的数据（$V / 4$ 字节）。最窄的对分垂直切过最长轴，穿过 $2 \cdot N / \max(A, B, …)$ 条链路（两个切面，计入环绕连接），因此单向带宽为 $N \cdot W_\text{ici} / \max(A, B, …)$。用数据量除以带宽，便得到上面的公式。

进一步了解 ReduceScatter

ReduceScatter 比初看起来更加基础，因为它实际上是 AllGather 的导数，反之亦然。也就是说，如果在前向传播中有：

那么我们会对反向模式导数 A’（通常在各分片上并不相同）执行 ReduceScatter，从而得到分片后的 A’：

同样，前向传播中的 $\text{ReduceScatter}_X(A[I] \{U_X\}) \to A[I_X]$ 意味着，反向传播中会有 $\text{AllGather}_{X}(A'[I_X]) \to A'[I]$。

把 AllReduce 拆成 AllGather 和 ReduceScatter 还有一个方便的性质：我们可以把最后一次 AllGather 推迟到更晚。很多时候，我们并不想付出在所有设备上重新组装完整矩阵乘积副本的成本。相反，即使在两个乘数的收缩维度都已分片的情况下，我们也希望保持分片状态：

在这种情况下，我们也可以执行 ReduceScatter 而不是 AllReduce，并选择在更晚的时候再执行 AllGather，即

请注意，ReduceScatter 会引入一个分片维度，因此在此处自然可以选择沿命名维度 I 或 K 进行分片。使用 ReduceScatter 时，我们通常需要选择要在哪个命名维度上引入新分片（不过，更大的建模上下文通常会迫使我们做出唯一选择）。正因如此，我们使用 ReduceScatter_X,K 语法来指定要分片的轴。

如何重叠矩阵乘法中的通信与计算

正如我们在第 1 部分中讨论的，只要通信足够快，我们通常假设总能让通信与某些有用计算重叠。本节中的集合通信通常可以与矩阵乘法计算本身重叠，但这样做并不简单。我们使用的算法称为集合通信矩阵乘法（collective matmul），最早由 Wang 等人描述。下面是一段简化动画，展示了如何实现这种重叠：

简而言之，我们可以在对矩阵的一个数据块执行矩阵乘法时，开始对之前的数据块执行环形归约。在某些情况下，还可以沿批次维度或矩阵输入维度分块。我们会在第 10 部分推导一个简单的 JAX 实现，Mosaic 文档也提供了一个很好的 GPU 示例。我们鼓励你有机会自己实现一个版本。

我们学到了什么？

• 数组的分片由 Mesh 和 Sharding 共同指定：Mesh 为 TPU 网格的物理硬件轴命名，Sharding 则把网格轴名称指派给数组的逻辑轴。

  • 例如，A[I_XY, J] 描述了一个抽象数组 A，其第一维沿 X 和 Y 两条网格轴分片。再结合 Mesh(mesh_shape=(4, 8), axis_names=(‘X’, ‘Y’))，或缩写形式 Mesh({‘X’: 4, ‘Y’: 8})，我们就知道该数组的第一维被分成了 32 份。

• 对分片数组执行算术运算与对未分片数组执行完全相同，除非你沿分片轴执行收缩。 在后一种情况下，必须引入某些通信。我们考虑了四种情形：

  1. 两个数组的收缩维度都未分片：不需要通信。
  2. 一个数组的收缩维度已分片（或者两个收缩维度沿不同轴分片）：执行运算前，我们先对其中一个输入执行 AllGather。
  3. 两个数组以相同方式沿收缩维度分片： 我们在局部把分片相乘，然后执行 AllReduce 或 ReduceScatter。
  4. 两个数组各有一个非收缩维度沿同一条网格轴分片： 我们先对其中一个输入执行 AllGather。

• TPU 大致使用 4 种核心通信原语：

  1. AllGather：$[A_X, B] \to [A, B]$
  2. ReduceScatter：$[A, B] \\{U_X\\} \to [A_X, B]$
  3. AllToAll：$[A, B_X] \to [A_X, B]$
  4. AllReduce：$[A_X, B]\\{U_Y\\} \to [A_X, B]$（严格来说不算原语，因为它把 ReduceScatter 与 AllGather 组合在一起）

• 这些操作各自的成本与延迟都不取决于轴的大小（只要它们受带宽限制），而只取决于输入数组的大小和链路带宽。对于单向 AllGather/ReduceScatter：

• AllReduce 由一次 ReduceScatter 后接一次 AllGather 构成，因此成本是上式的 2 倍。AllToAll 只需让分片沿环传递部分路程，因此其成本是 AllGather 的 ¼。总结如下：

练习题

下面是一些基于本节内容、很有启发性的练习。暂时不会给出全部答案，但我们会尽可能继续补充答案。

问题 1［复制式分片］： 一个数组的分片为 $A[I_X, J, K, \ldots]$（即只沿 $X$ 分片），网格为 Mesh({'X': 4, 'Y': 8, 'Z': 2})。在所有芯片上，$A$ 占用的总字节数与一份数组副本的大小之比是多少？

问题 2［AllGather 延迟］： 在网格为 Mesh({'X': 4, 'Y': 4, 'Z': 4}) 的 TPU v4p 4x4x4 切片上，如果数据类型为 bfloat16，$\text{AllGather}_X([B_X, D_Y])$ 需要多长时间？这里 $B=1024$、$D=4096$。$\text{AllGather}_{XY}([B_X, D_Y])$ 呢？$\text{AllReduce}_Z([B_X, D_Y] \{U_Z \})$ 呢？

问题 3［延迟受限的 AllGather］： 假设我们正在执行 $\text{AllGather}_X([B_X])$，但 $B$ 很小（比如 128）。在网格为 Mesh({'X': 4, 'Y': 4, 'Z': 4}) 的 TPU v4p 4x4x4 切片上，以 bfloat16 格式执行需要多长时间？提示：你很可能受延迟限制。

问题 4［矩阵乘法策略］： 为执行 $X[B, D] \cdot_D Y[D_X, F] \to Z[B, F]$，本节建议先执行 $\text{AllGather}_X(Y[D_X, F])$，再对完全复制的矩阵相乘（情形 2，策略 1）。另一种做法是像 $X[B, D_X] \cdot_D Y[D_X, F] \to Z[B, F] \\{U_X\\}$ 那样对局部分片相乘（情形 3，策略 2），然后执行 $\text{AllReduce}_X(Z[B, F] \\{ U_X\\})$。两种方法各自执行多少 FLOP 和通信？哪一种更好，为什么？

问题 5［最低延迟］： 假设我想在 TPU v4p 4x4x4 上以尽可能低的延迟执行矩阵乘法 $A[I, J] \cdot_J B[J, K] \to C[I, K]$。假设输入可以任意分片，但结果应当完全复制。输入应如何分片？FLOP 总数和通信时间分别是多少？

问题 6： 假设我们想在 TPU v5e 4x4 上执行 $A[I_X, J_Y] \cdot_J B[J_Y, K] \to C[I_X, K]$。需要执行什么通信？通信与计算各花费多少时间？

• 那么 $A[I_X, J] \cdot_J B[J_X, K_Y] \to C[I_X, K_Y]$ 呢？这是训练中最标准的场景，我们在其中组合数据分片、张量分片与 ZeRO 分片。
• 那么 $A[I_X, J] \cdot_J B[J, K_Y] \to C[I_X, K_Y]$ 呢？这是推理中的标准场景，其中执行纯张量并行（+数据并行）。

问题 7： 一个典型的 Transformer 块包含两个矩阵 $W_\text{in}[D, F]$ 和 $W_\text{out}[F, D]$，其中 $F \gg D$。假设批大小为 B，那么整个块就是 $In[B, D] \cdot W_\text{in}[D, F] \cdot W_\text{out}[F, D]$。取 $D=8192$、$F=32768$、$B=128$，并假设所有数据都采用 bfloat16。假设在 TPU v5e 2x2 切片上运行，但姑且设每块 TPU 只有 300MB 可用内存。应如何对 In、$W_\text{in}$、$W_\text{out}$ 和 Out 分片，才能既不超过内存限制，又尽可能缩短总时间？通信和 FLOP 分别花费多长时间？提示：最终输出不必完全复制，但它应当与输入采用相同的分片方案，以便重复应用这个“层”。

问题 8［挑战］： 以上面的简短代码片段为模板，分配一个分片数组，并使用 pmap 或 shard_map 对 4 种主要通信原语（AllGather、AllReduce、ReduceScatter 和 AllToAll）逐一进行基准测试。你需要使用 jax.lax.all_gather、jax.lax.psum、jax.lax.psum_scatter 和 jax.lax.all_to_all。你理解这些函数的语义吗？它们各自需要多长时间？

问题 9［分片矩阵乘法的另一种策略？］： 我们在上文声称，当矩阵乘法中只有一个输入沿收缩维度分片时，应当对分片矩阵执行 AllGather，然后在局部执行所得收缩。你或许会想到另一种策略：执行分片矩阵乘法，然后对结果执行 AllReduce（仿佛两个输入都沿收缩维度分片），也就是通过下面两步计算 $A[I, J_X] *_J B[J, K] \to C[I, K]$：

1. $C[I, K] \\{ U_X \\} = A[I, J_X] \cdot B[J_X, K]$
2. $C[I, K] = \text{AllReduce}(C[I, K] \\{ U_X\\})$

请回答：

1. 对矩阵 $A[N, M]$ 和 $B[M, K]$ 明确写出该算法，使用索引准确展示在哪台设备上执行什么计算。假设 $A$ 以 $A[I, J_X]$ 的方式分片到 ND 台设备上，并希望输出复制到所有设备。
2. 现在假设你可以接受最终结果不在每台设备上复制，而是分片（沿 N 或 K 维度均可）。上述算法应如何改变？
3. 只考虑上述策略（第 2 问中的策略，而不是第 1 问）的通信成本，它与先对 A 执行 AllGather、再执行矩阵乘法的算法相比如何？

问题 10：AllToAll 趣题： 上表指出，在吞吐量受限的区域内，执行 AllToAll 的时间比执行 AllGather 或 ReduceScatter 低 4 倍。在这道题中，我们将理解这个 4 倍因子从何而来，也会看看如果只有单向 ICI 链路而非双向 ICI 链路，这个因子会如何变化。

1. 先从单向情况开始。设有 D 台设备组成环形拓扑，并希望对 N x N 矩阵 $A[I_X, J]$ 执行 AllGather 或 ReduceScatter（为简单起见，假设 $D$ 整除 $N$）。描述这两种集合通信涉及的通信，并计算整个算法期间通过单条 ICI 链路传输的标量（浮点数或整数）总数。
2. 接下来考虑 AllToAll，仍然采用单向 ICI。此时的算法与 AllGather 有何不同？计算该算法中通过单条 ICI 链路传输的标量数。
3. 你应该会发现，第 (a) 问与第 (b) 问答案之比是个漂亮的数字。请用简单语言解释这个因子从何而来。
4. 现在加入双向通信。这对 AllGather 所需总时间有何影响？
5. 加入双向通信对 AllToAll 所需总时间有何影响？
6. 现在请直接解释双向环中 AllGather 时间与 AllToAll 时间的比值。

第 3 部分到此结束！要继续阅读第 4 部分（关于 Transformer 数学），请点击这里！