集体矩阵乘法

集合矩阵乘法#

张量并行 (TP) 和数据并行 (DP) 是最常用的并行技术，它们使得在多个加速器上运行越来越大的模型成为可能。然而，联合使用这些技术意味着在程序中，我们有时会得到以无法直接执行操作的方式分片的数据，除非进行额外的通信。这种问题常见于 Transformer 的 MLP 块的开头。在那里，输入激活可能在批处理轴上进行分片 (DP)，而权重可能在输出特征维度上进行分区 (TP)。

Left matrix is split into halves by rows, right matrix is split into halves by columns

收缩维度未被分片，因此看起来我们似乎可以直接相乘输入，但存在一个问题：输出无法在其两个维度上同时沿同一设备轴进行分片！

有一个简单的方法可以解决这个问题：我们可以对激活或权重执行 All-Gather（此处我们侧重于激活侧），然后与另一个分片的操作数执行局部矩阵乘法。这个简单的策略有效，但有一个缺点：在 All-Gather 运行期间我们无法开始计算矩阵乘法！这意味着我们的硬件利用率不足！

为了实现更高的利用率，我们将展示实现 Pallas:MGPU 内核是多么简单，该内核将跨设备通信与矩阵乘法重叠，在足够大的问题规模下实现近乎最优的利用率。我们的实现大量使用了 NVLINK 互连，这允许我们在不涉及主机的情况下执行高带宽的 GPU 间通信。

这种方法已经产生了显著的性能提升！如果我们考虑一个 M=1024、K=4096 和 N=4096 的 f16 矩阵乘法以及正态分布数据，我们的基准测试表明在单个 H100 上大约需要 43us。在下表中，我们扩大了 M 维度，使得每个分片的形状为 M=1024。我们可以通过将局部运行时间估计值乘以设备数量，并为每一轮通信增加约 6us（与同步相关的内存屏障开销很大），来计算分布式内核执行的预期下限。对我们的内核进行基准测试得出以下结果

设备数量

内核时间

TC 利用率

理论下限

TC 利用率

参考时间

TC 利用率

2

102us

68%

92us

75%

147us

47%

4

212us

66%

190us

73%

290us

48%

8

436us

64%

386us

72%

565us

49%

正如您所见，这里仍有一些优化空间，但与使用 NCCL All-Gather 和 cuBLAS 矩阵乘法的基准实现相比，我们至少获得了更好的利用率。

算法概述：环形 All-Gather#

为了计算 AllGather(A) @ B，我们在参与计算的 D 个设备上形成一个环。在每个步骤中，设备获取上一次接收到的分片（从其本地分片开始），并将其传递给环中的下一个设备。在发送发生的同时，我们计算上一次接收到的 A 分片与本地 B 分片之间的矩阵乘法。

all_gather

更正式地说，该算法分 D 个步骤进行。在步骤 i (0 <= i < D) 中，设备 d 从设备 (d + 1) % D 接收分片 A_{(d + i) % D}（实际上第一步不接收），计算 A_{(d + i) % D} @ B_d，并将结果写入输出缓冲区的一个切片。在进行计算的同时，设备 d 将分片 A_{(i + d) % D} 发送给设备 (i - 1) % D，以便其在步骤 i + 1 中使用（最后一步不发送）。经过 D 个步骤后，设备 d 将看到 A 的每个分片并计算出完整的输出。

用于跨设备通信的 Pallas 原语#

我们使用三个 Pallas 函数进行跨设备通信

  • plgpu.remote_ref(ref, device_id): 此函数获取全局内存 (GMEM) 中缓冲区的引用，并返回一个指向由 device_id 指定的不同设备上相同缓冲区的引用。通过 NVLINK 通信时，即使其数据位于远程内存中，也可以直接读取或写入此引用。

  • pl.semaphore_signal(sem, device_id=...): 增加目标设备上的信号量。这通常用于指示某个过程的完成，例如当我们通知远程设备它正在等待的数据已被发送时。

  • pl.semaphore_wait(sem, value=..., decrement=...): 阻塞直到本地信号量达到特定值。如果 decrement 为 True（默认），信号量的值将减少等待的数量。如果为 False，操作效率更高，但在等待完成后不会修改信号量的值。这常用于等待来自远程设备的信号。

使用 Pallas 进行实现#

注意

在此，我们仅展示内核的简化版本，这使我们能够专注于最有趣的细节。您可以在我们的示例目录中找到完整的实现。

首先，我们关注内核的设置。对于计算部分，我们将重用 hopper_matmul_mgpu 中优化的矩阵乘法内核实现。由于计算内核将利用 warp 特化，我们使用 3 个 Pallas 线程。它也是持久的，这意味着我们启动一个与 SM 数量一样大的网格（从 JAX 设备上的 .core_count 查询得到）。计算内核使用 pl.run_scoped 进行 SMEM 分配，因此我们不使用 scratch_shapes。

def all_gather_lhs_matmul(
    lhs: jax.Array,
    rhs: jax.Array,
    axis_name,
    *,
    config: hopper_matmul_mgpu.TuningConfig,
    dtype: jnp.dtype = jnp.bfloat16,
) -> jax.Array:
  if (num_devices := jax.device_count()) != jax.process_count():
    raise ValueError("The kernel only supports one device per process")
  if (axis_size := lax.axis_size(axis_name)) != num_devices:
    raise ValueError("The kernel can only work over all devices in a Mesh.")
  ...

  m_shard, k = lhs.shape
  _, n_shard = rhs.shape
  tile_m, tile_n, tile_k = config.tile_m, config.tile_n, config.tile_k
  cta_tile_m = tile_m * (1 + (config.wg_dimension == MatmulDimension.M))
  num_sms = jax.extend.backend.get_default_device().core_count

  def kernel_body(lhs_local_ref, rhs_ref, out_ref, scratch_ref):
    ...

  result, _ = plgpu.kernel(
      kernel_body,
      out_shape=[
          # The output (with M gathered)
          jax.ShapeDtypeStruct((axis_size * m_shard, n_shard), dtype),
          # A scratch buffer for LHS all-gather
          jax.ShapeDtypeStruct((axis_size - 1, m_shard, k), dtype),
      ],
      grid=(num_sms,),
      num_threads=3, # The matmul kernel uses 3 threads: 2 compute and 1 memory
      thread_name="wg",
  )(lhs, rhs)
  return result

上面的内核有两个输出。第一个是我们原语的实际结果，第二个用作接收左操作数的暂存空间。请注意，我们可以将引导轴缩小到小于 axis_size - 1，但在那种情况下，我们需要向发送设备引入反压，这需要额外的昂贵通信。

注意

您可以查看 TPU 分布式通信指南，了解如何处理这种反压。

现在让我们看看内核主体的轮廓

def all_gather_lhs_matmul(...):
  def kernel_body(lhs_local_ref, rhs_ref, out_ref, scratch_ref, out_smem, received_sem):
    wg_idx = lax.axis_index("wg")
    dev_id = lax.axis_index(axis_name)
    # This device sends to dev_id - 1, forming a ring.
    send_dev_id = lax.rem(dev_id + axis_size - 1, axis_size)
    send_scratch_ref = plgpu.remote_ref(scratch_ref, send_dev_id)

    def device_step(lhs_source_ref, device_offset):
      # Invariant: lhs_source_ref contains A_{(dev_id + device_offset) % D}
      # and is ready to be used for computation.

      ...

    # We peel the first step to read data directly from lhs_local_ref.
    device_step(lhs_local_ref, 0)
    @pl.loop(1, num_devices)
    def _device_loop(device_offset):
      device_step(scratch_ref.at[device_offset - 1], device_offset)

我们通过查询 lax.axis_index(axis_name) 来确定我们在环中的位置，并计算我们将向其发送数据的下一个设备的索引 (send_dev_id)。然后，我们循环调用 device_body，次数与设备数量相同。我们剥离循环的第一步，因为我们仅在该步骤中使用本地引用作为发送源（之后发送源自之前在暂存缓冲区中接收到的数据）。

我们现在准备研究主循环

def all_gather_lhs_matmul(...):
  ...

  def kernel_body(lhs_local_ref, rhs_ref, out_ref, scratch_ref, out_smem, received_sem):
    ...

    def device_step(lhs_source_ref, device_offset):
      # We are computing block (dev_id + device_offset) % D of the output.
      out_device_idx = lax.rem(device_offset + dev_id, axis_size)
      out_device_m_slice = pl.ds(out_device_idx * m_shard, m_shard)

      # In step `device_offset`, we send A_{(dev_id + device_offset) % D} to
      # the next device in the ring, into scratch slot `device_offset`.
      # We also don't send on the last step since that would return the data
      # back to its original source.
      next_scratch_slot = device_offset
      is_send_wg = wg_idx == 0 # Only one warpgroup per CTA sends
      has_send_space = next_scratch_slot < axis_size - 1
      should_send = is_send_wg & has_send_space

      # This function will be called by hopper_matmul_mgpu.kernel in the body
      # of its pipeline. We use it to take the tile of LHS loaded into SMEM and
      # issue a TMA send to the next device in the ring.
      def send_lhs(m_idx, n_idx, k_idx, a_smem, b_smem, send_ref, should_send):
        del b_smem  # Unused.
        # We only send when n_idx == 0 to avoid sending the same data
        # multiple times when revisiting the left operand.
        @pl.when(should_send & jnp.bool(n_idx == 0))
        def _():
          k_slice = pl.ds(k_idx * tile_k, tile_k)
          m_slice = pl.ds(m_idx * cta_tile_m, cta_tile_m)
          plgpu.copy_smem_to_gmem(a_smem, send_ref.at[m_slice, k_slice])
          # Wait for previous copies to complete. We pass in delay_release=1
          # to the pipeline in the matmul kernel to ensure that it doesn't
          # overwrite the input until at least the next step completes, but it
          # will not wait any longer.
          plgpu.wait_smem_to_gmem(1, wait_read_only=True)

      hopper_matmul_mgpu.kernel(
          lhs_source_ref,  # LHS shard for this step
          rhs_ref,  # RHS shard is always the same
          out_ref.at[out_device_m_slice],  # Slice of output to update
          out_smem,
          config=config,
          pipeline_callback=functools.partial(
              send_lhs,
              send_ref=send_scratch_ref.at[next_scratch_slot],
              should_send=should_send,
          ),
          delay_release=1,
      )

      # Wait for the next scratch to arrive for the next step's computation.
      # Each device signals its neighbor when it has finished sending.
      @pl.when(should_send)
      def _signal():
        # Make sure our remote copy is done, then signal.
        plgpu.wait_smem_to_gmem(0, wait_read_only=False)
        pl.semaphore_signal(received_sem, device_id=send_dev_id)
      @pl.when(has_send_space)
      def _wait():
        # Here, we wait for the data to arrive from the previous device in the
        # ring. At each step, will expect to receive a signal from each SM.
        # We use decrement=False to make this operation slightly faster, but
        # this also means that we need to scale the expected number of signals
        # by the number of steps taken so far (as the value only increases).
        pl.semaphore_wait(received_sem, value=(device_offset + 1) * num_sms, decrement=False)

    ...

这里按顺序发生了一些事情

  1. 我们首先计算在循环此步骤中将计算的输出切片。

  2. 然后，我们调用优化的矩阵乘法内核，并注入一个 pipeline_callback。我们利用计算内核必须将左操作数提取到 SMEM 中的事实，并指示 TMA 引擎将本地数据异步流式传输到下一个设备。流量由硬件透明地通过 NVLINK 路由。值得注意的是，我们仅从其中一个计算线程发出发送指令，并且仅在我们第一次访问左操作数时（为了计算许多输出切片，它可能会被多次重新加载）。

  3. 最后，发送线程确保发送已完成，并向接收设备上的 received_sem 发出信号以表明这一点。在此之后，所有线程都会等待，直到确定已接收到循环下一步所需的所有数据（最后一步会跳过等待）。

将内核集成到 JAX 中#

要调用内核，您需要将其包装在 jax.shard_map 中

m_shard, n_shard, k = 1024, 1024, 1024
dtype = jnp.float16
mesh = jax.make_mesh((jax.device_count(),), ("x",),
                     axis_types=(jax.sharding.AxisType.Explicit,))
with jax.set_mesh(mesh):
  a = jax.random.normal(jax.random.key(1), (m_shard * jax.device_count(), k), dtype)
  b = jax.random.normal(jax.random.key(2), (k, n_shard * jax.device_count()), dtype)
  a = jax.sharding.reshard(a, P("x", None))
  b = jax.sharding.reshard(b, P(None, "x"))

  # Example config for 8xH100. You might need to retune to your shape.
  config = hopper_matmul_mgpu.TuningConfig(
      tile_m=128, tile_n=128, tile_k=64, max_concurrent_steps=4,
      grid_minor_dim=MatmulDimension.N, grid_tile_width=8,
      wg_dimension=MatmulDimension.N,
  )

  kernel = jax.jit(
      jax.shard_map(
          functools.partial(all_gather_lhs_matmul, axis_name="x", config=config),
          out_specs=P(None, "x"),
          check_vma=False,
      )
  )
  c = kernel(a, b)