10 TPU 与分布式通信原语

最新推荐文章于 2026-08-18 16:57:45 发布

原创

最新推荐文章于 2026-08-18 16:57:45 发布 · 337 阅读

之前的内容都是在单张 GPU 上展开的，但实际 LLM 时代，分布式训练/推理已经成为绕不开的话题。本节以 TPU 为切入点，介绍 TPU 在分布式场景下常用的分布式通信原语实现，以及其应用。其思路和实现原理同样适用于 GPU。

TPU 的硬件模型

TPU 通常的物理拓扑是 mesh。比如 2*2 的mesh，组成一个ring。

TPU 每条链路单向带宽大约为 92 GB/s，单 hop 延迟大约 1 微秒。

每个 TPU v5p 设备有两个 core，每个 core 是单线程的 VLIW 风格处理器，但 SIMD 很宽：对 fp32 而言，原生处理宽度是 1024 元素；TPU通常把它组织成一个 (8, 128) 的 tile。 Pallas 对齐、切片、RDMA 连续性基本都围着这个最小高效单位在转。

TPU 的内存层级和 GPU 很不一样。TPU 最重要的是 VMEM，也就是每个 core 拥有的 64 MiB 高带宽 SRAM scratchpad。与 H100 相比，这个 scratchpad 容量非常大。除此之外，每个 device 还有 95 GB 的 HBM2e；但和 GPU 复杂的 cache hierarchy 相比，TPU 更像“HBM + 大片上 scratchpad”的体系。很多实现策略，本质上都是在问：哪些数据值得常驻 VMEM，哪些访问必须走 RDMA，哪些布局才能让编译器和 ICI 都满意。

TPU Pallas

Pallas 是一个很低层、很贴近硬件的系统，但它采用 tracing-based compilation model，所以它只是“发射指令的模板”，不是直接在 TPU 上跑的动态程序。

我觉得最值得记住的是下面四句话。

Python control flow 不是 kernel control flow。 在 traced kernel 里写 if、while，本质发生在 trace time，而不是 runtime。只有用 pl.when、lax.cond、lax.fori_loop 之类 traced 控制流，才是真正把分支或循环发到 TPU 上执行。
Python print 不是 kernel print。 想看 runtime 的动态值，必须用 jax.debug.print，而且还要打开 starter code 里的 debug 开关。
性能只取决于最终发射出的指令序列。 你在 Python 里如何组织代码、是否用了字典或 helper function，并不直接决定 runtime 性能；真正决定性能的是 tracing 之后的底层指令。
Array ref 和 array value 要分清。 ref 更像带 shape 信息的指针，value 才是被临时放进向量寄存器里的数值。RDMA 操作的是 ref，数学运算操作的是 value。

TPU RDMA

Pallas 会保证各设备上 kernel address space 布局一致，所以“我本地这个 buffer 的 ref”也可以用来指向远端设备上对应位置的 buffer。

RDMA 是异步的，因此必须配套 semaphore 和 wait 语义来管理生命周期。

示例 TPU RDMA 代码：

import time

import jax
import numpy as np
from jax import lax, numpy as jnp
from jax.experimental import pallas as pl
from jax.experimental.pallas import tpu as pltpu
from jax.sharding import Mesh, PartitionSpec

AXIS_NAME = "i"
N_DEVICES = 4

ENABLE_DEBUG = False


def pallas_get_my_device_id():
    return lax.axis_index(AXIS_NAME)


def pallas_rdma_start(*, src_ref, dst_ref, dst_device_id, src_send_sem, dst_recv_sem):
    pltpu.make_async_remote_copy(
        src_ref=src_ref,
        dst_ref=dst_ref,
        send_sem=src_send_sem,
        recv_sem=dst_recv_sem,
        device_id=dst_device_id,
        device_id_type=pltpu.DeviceIdType.LOGICAL,
    ).start()


def pallas_rdma_wait_send(*, src_ref, src_send_sem):
    pltpu.make_async_remote_copy(
        src_ref=src_ref,
        dst_ref=src_ref,  # ignored by 'wait_send'
        send_sem=src_send_sem,
        recv_sem=src_send_sem,  # ignored by 'wait_send'
        device_id=0,  # ignored by 'wait_send'
        device_id_type=pltpu.DeviceIdType.LOGICAL,
    ).wait_send()


def pallas_rdma_wait_recv(*, dst_ref, dst_recv_sem):
    pltpu.make_async_remote_copy(
        src_ref=dst_ref,  # ignored by 'wait_recv'
        dst_ref=dst_ref,
        send_sem=dst_recv_sem,  # ignored by 'wait_recv'
        recv_sem=dst_recv_sem,
        device_id=0,  # ignored by 'wait_recv'
        device_id_type=pltpu.DeviceIdType.LOGICAL,
    ).wait_recv()


def exchange_with_neighbor_pallas_scratch_specs(x):
    return {
   
   
        "send_sem": pltpu.SemaphoreType.DMA,
        "recv_sem": pltpu.SemaphoreType.DMA,
    }


def exchange_with_neighbor_pallas_kernel(x_ref, out_ref, scratch_refs):
    my_device_id = pallas_get_my_device_id()
    paired_device_id = my_device_id + 1 - 2 * lax.rem(my_device_id, 2)

    send_sem = scratch_refs["send_sem"]
    recv_sem = scratch_refs["recv_sem"]

    pallas_rdma_start(
        src_ref=x_ref,
        dst_ref=out_ref,
        dst_device_id=paired_device_id,
        src_send_sem=send_sem,
        dst_recv_sem=recv_sem,
    )
    pallas_rdma_wait_send(src_ref=x_ref, src_send_sem=send_sem)
    pallas_rdma_wait_recv(dst_ref=out_ref, dst_recv_sem=recv_sem)

reduce-scatter 与 all-gather

reduce-scatter：“先全局求和，再把结果按 device 切开”；
all-gather：“先把各 device 的 shard 拼起来，再在每个 device 上都放一份完整结果”。

如果只说语义，这两个 collective 很简单；真正的难点是怎样在 ring 上把它们写成高吞吐的 RDMA schedule。对于 4 个设备的 ring，实现思路可以概括成下面两句：

all-gather：每个设备先把自己的 shard 写到输出中对应的位置，然后不断把“已经拥有的一块”发给下一个设备，同时从上一个设备收一块，直到拼齐完整结果。
reduce-scatter：每个设备先把某一块发出去、收回来某一块 partial，然后把收到的块和自己本地对应块相加，再把新的 partial sum 继续往 ring 里传，最后留下属于自己的那一块。

如果把 ring 想成一个流动的管道，那么 all-gather 传播的是“原始 chunk”，reduce-scatter 传播的是“不断累加的 partial chunk”。两者看上去很像，但一个的 payload 是原始值，一个的 payload 是归约中的中间值。

参考reduce-scatter 与 all-gather实现：

def reduce_scatter_pallas_scratch_specs(x):
    shard_shape = (x.shape[0] // N_DEVICES, x.shape[1], x.shape[2])
    return {
   
   
        "carry_buf": pltpu.VMEM(shape=shard_shape, dtype=x.dtype),
        "recv_buf": pltpu.VMEM(shape=shard_shape, dtype=x.dtype),
        "send_sem": pltpu.SemaphoreType.DMA,
        "recv_sem": pltpu.SemaphoreType.DMA,
    }


def reduce_scatter_pallas_kernel(x_ref, out_ref, scratch_refs):
    my_device_id = pallas_get_my_device_id(