[Pallas:MGPU] Add a collective matmul tutorial

docs/pallas/gpu/collective_matmul.md

@@ -0,0 +1,305 @@
+# Collective matrix multiplication
+
+Tensor parallelism (TP) and data parallelism (DP) are the most frequently used
+parallelism techniques that make it possible to fit the ever larger models onto
+a number of accelerators. However, their joint use means that in our programs,
+we sometimes end up with data sharded in ways that don't make it directly
+possible to execute an operation without additional communication. One such
+problem frequently happens at the beginning of the MLP block of a Transformer.
+There, the input activations might be sharded on the batch axis (DP), while the
+weights might be partitioned on the output feature dimension (TP).
+
+<center><img src="../../_static/pallas/distributed/collective_all_gather_operands.svg" alt="Left matrix is split into halves by rows, right matrix is split into halves by columns"></center>
+
+The contraction dimension is not sharded, so it might seem that we can just
+multiply the inputs, but there is a problem: the output can't be sharded along
+the same device axis on both of its dimensions!
+
+There's a simple way to solve this problem: we can all-gather activations or
+weights (here we focus on the activation side), and then perform a local matrix
+multiplication with the other operand sharded. This simple strategy works, but
+it has a downside: we can't begin computing the matrix multiplication while the
+all-gather is running! That means we're underutilizing our hardware!
+
+To achieve better utilization, we'll show how simple it is to implement a
+Pallas:MGPU kernel that overlaps the cross-device communication with the
+matrix-multiplication, achieving almost optimal utilization on large enough
+problem shapes. Our implementation makes heavy use of the NVLINK interconnect,
+which allows us to perform high-bandwidth inter-GPU communication without
+involving the host.
+
+This approach already yields considerable performance improvements! If we
+consider a f16 matmul with M=1024, K=4096 and N=4096 and normally distributed
+data, our benchmarks indicate that it should take about 43us on a single H100.
+In the table below, we scale up the M dimension so that the per-shard shape is
+M=1024. We can compute an expected lower bound for the execution of our
+distributed kernel by multiplying that local runtime estimate by the number of
+devices and by adding about 6us for each round of communication (the memory
+fences associated with the synchronization are expensive). Benchmarking our
+kernel yields the following results:
+
+| Device count | Kernel time | TC utilization | Lower bound | TC utilization | Reference time | TC utilization |
+|--------------|-------------|----------------|-------------|----------------|----------------|----------------|
+| 2            | 102us       | 68%            | 92us        | 75%            | 147us          | 47%            |
+| 4            | 212us       | 66%            | 190us       | 73%            | 290us          | 48%            |
+| 8            | 436us       | 64%            | 386us       | 72%            | 565us          | 49%            |
+
+As you can see there are still some opportunities for optimization here, but at
+least we're getting much better utilization compared to the baseline
+implementation of a NCCL all gather and cuBLAS matmul.
+
+## Algorithm overview: Ring All-Gather
+
+To compute `AllGather(A) @ B`, we form a ring on the participating `D` devices.
+At each step, the device takes the last received shard (starting from its local
+shard), and passes it to the next device in the ring. While the send is
+happening, we compute the matrix multiplication between the last received `A` shard
+and the local `B` shard.
+
+![all_gather](../../_static/pallas/distributed/all_gather.svg)
+
+More formally, the algorithm proceeds in `D` steps. In step `i` (`0 <= i < D`),
+device `d` receives shard `A_{(d + i) % D}` (we don't actually receive in the
+first step) from device `(d + 1) % D`, computes `A_{(d + i) % D} @ B_d`, and
+writes the result to a slice of the output buffer.  Concurrently with the
+compute, the device `d` sends shard `A_{(i + d) % D}` to device `(i - 1) % D`
+for its use in step `i + 1` (we don't send in the last step). After `D` steps,
+device `d` will have seen every shard of `A` and computed the full output.
+
+## Pallas primitives for inter-device communication
+
+We use three Pallas functions for inter-device communication:
+
+*   **`plgpu.remote_ref(ref, device_id)`**: This function takes a reference to a
+    buffer in global memory (GMEM) and returns a reference to the same buffer on a
+    *different* device, specified by `device_id`. When communicating over NVLINK,
+    this reference can be read or written to directly, even though its data is located
+    in remote memory.
+*   **`pl.semaphore_signal(sem, device_id=...)`**: Increments a semaphore on a
+    target device. This is usually used to indicate completion of some process,
+    such as when we notify the remote device that the data it's waiting for has
+    been sent.
+*   **`pl.semaphore_wait(sem, value=..., decrement=...)`**: Blocks until a local
+    semaphore reaches a certain value. If decrement is `True` (default), the
+    value of the semaphore is decreased by the awaited amount. If it is `False`,
+    the operation is more efficient, but it does not modify the value of the
+    semaphore after the wait completes. This is frequently used to await signals
+    from a remote device.
+
+## Implementation with Pallas
+
+```{note}
+Here, we only present a simplified version of the kernel, which allows us to
+focus on the most interesting details. You can find [the full implementation in
+our examples directory](https://github.com/jax-ml/jax/blob/main/jax/experimental/pallas/ops/gpu/collective_matmul_mgpu.py).
+```
+
+First, we focus on the set-up of our kernel. For the compute part, we will reuse
+our optimized matmul kernel implementation from `hopper_matmul_mgpu`. Since the
+compute kernel will utilize warp-specialization, we use 3 Pallas threads. It
+is also persistent, which means that we launch a grid as large as the number of
+SMs (queried from `.core_count` on the JAX device). The compute kernel uses
+`pl.run_scoped` for SMEM allocations, so we don't use `scratch_shapes`.
+
+```python
+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
+```
+
+The kernel above has two outputs. First one is the actual result of our
+primitive, while the second one is used as a scratch space to receive the left
+operands. Note that we could shrink the leading axis to be smaller than
+`axis_size - 1`, but at that point we would need to introduce backpressure to
+the sending devices, which requires additional expensive communication.
+
+```{note}
+You can see how to deal with this backpressure in the [TPU distributed communication guide](../tpu/distributed.md#run-ahead-and-race-conditions).
+```
+
+Let us now look at the outline of the kernel body:
+
+```python
+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, device_id_type=pl.DeviceIdType.LOGICAL
+    )
+
+    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)
+```
+
+We locate our position in the ring by querying `lax.axis_index(axis_name)` and
+compute the index of the next device, to which we will be sending the data
+(`send_dev_id`). Then, we loop over invocations of the `device_body` as many
+times as there are devices. We peel the first step of the loop, because we use
+the local reference as the source for the send in that step only (after that the
+sends originate from the data previously received in the scratch buffer).
+
+We are ready to investigate the main loop now:
+
+```python
+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)
+
+    ...
+```
+
+A few things happen here in a sequence:
+1. We begin by computing the slice of the
+  output that we will compute at this step of the loop.
+2. Then, we call into the optimized matmul kernel, but injecting it with a
+  `pipeline_callback`. We use it to take advantage of the fact that the compute
+  kernel has to fetch the left operand into SMEM, and we instruct the TMA engine
+  to asynchronously stream the local data to the next device. The traffic is
+  transparently routed through NVLINK by the hardware. It is worth noting that we
+  only issue sends from one of the compute threads and only when we visit the left
+  operand for the first time (it might be reloaded many times to compute many
+  output tiles).
+3. Finally, the sending thread makes sure that the sends have completed and
+   signals the `received_sem` on the receiving device to indicate that. After
+   that, all threads wait until they are sure that all the data for the next
+   step of the loop has been received (the wait is skipped on the last step).
+
+## Integrating the kernel with JAX
+
+To invoke the kernel, you need to wrap it into `jax.shard_map`:
+```python
+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)
+```

docs/pallas/gpu/index.rst

@@ -9,6 +9,7 @@ Backend specific documentation for the Mosaic GPU backend.
    reference
    pipelining
    blackwell_matmul
+   collective_matmul
 
 .. toctree::
    :caption: Guides