JAX 分布式基础：从单卡到多卡的思维跃迁

说实话，我第一次接触 JAX 的分布式编程时，心里是有点抵触的。毕竟之前用惯了 PyTorch 的 DistributedDataParallel，总觉得分布式训练不就是把模型复制到多张卡上，同步一下梯度嘛。但 JAX 的哲学完全不同——它让你从「数据怎么分」这个层面开始思考，而不是事后补救。

这一章，我们就来聊聊 JAX 分布式最核心的几个概念。嗯，我会尽量用我踩过的坑来帮你理解。

pmap：最朴素的并行方式

pmap 是 JAX 里最直观的并行原语。它的作用很简单：把一个函数「映射」到多个设备上执行。

我习惯这么理解：pmap 就像工厂里的流水线，你把一批数据分成几份，每个工人（设备）处理一份。工人之间怎么通信？pmap 默认是不管的，除非你显式调用集体通信原语。

import jax
import jax.numpy as jnp

# 假设我们有 4 个设备
devices = jax.devices()
print(f"可用设备数: {len(devices)}")

# 定义一个简单的函数
def f(x):
    return x * 2

# 用 pmap 包装
parallel_f = jax.pmap(f)

# 输入数据：shape 为 (8, 3)
data = jnp.ones((8, 3))
result = parallel_f(data)

print(result.shape)  # 输出: (8, 3)
# 实际上每个设备处理了 2 个样本

这里有个坑，我曾经踩过：pmap 要求输入的第一个维度必须能被设备数整除。如果你有 4 张卡，batch size 是 10，那就会报错。为什么？因为 JAX 希望每个设备拿到相同数量的数据，这样计算图才是确定的。

shard_map：更精细的控制

shard_map 是 JAX 后来引入的，比 pmap 更灵活。我个人觉得，shard_map 才是真正适合生产环境的选择。

为什么？因为 pmap 只能沿着第一个轴切分，而 shard_map 允许你指定任意轴的切分方式。举个例子：

from jax.sharding import PartitionSpec as P
from jax.experimental.shard_map import shard_map

# 定义一个函数，输入是 (batch, seq_len, hidden)
def attention(x):
    # 假设这里做自注意力
    return x

# 只在 batch 维度切分
sharded_attn = shard_map(
    attention,
    mesh=mesh,
    in_specs=(P('batch', None, None),),
    out_specs=(P('batch', None, None),)
)

你看，in_specs 里的 P('batch', None, None) 表示：第一个维度按 'batch' 轴切分，后两个维度不切分（None）。这种表达能力，在模型并行（比如张量并行）时特别有用。

我记得有一次做 GPT 模型的张量并行，需要把注意力头的维度切到不同设备上。用 pmap 根本做不到，但 shard_map 一行代码就搞定了。

Sharded Device Array：你看不见的数据分布

当你用 pmap 或 shard_map 处理数据后，得到的是一个 ShardedDeviceArray。这东西长什么样？

# 接上面的例子
print(type(result))  # <class 'jax.interpreters.pxla.ShardedDeviceArray'>
print(result.device_buffers)  # 每个设备上的实际数据块

从用户角度看，它就像一个普通的 JAX 数组。你可以对它做切片、索引、数学运算。但底层，数据是分散的。当你调用 result[0] 时，JAX 会自动从对应的设备上拉取数据。

这里有个性能陷阱：如果你频繁地在不同设备间访问数据，通信开销会很大。我见过有人写训练循环时，每步都打印 loss 值，结果 loss.item() 触发了跨设备的数据聚合，训练速度直接掉了 30%。

集体通信原语：设备间的悄悄话

分布式训练离不开设备间的通信。JAX 提供了一套类似 MPI 的集体通信原语，但用起来更函数式。

原语	作用	我常用的场景
psum	所有设备求和	计算全局梯度
pmean	所有设备求平均	计算平均 loss
all_gather	收集所有设备的数据	全量梯度检查
pmax	所有设备求最大值	梯度裁剪中的全局 norm

这些原语必须在 pmap 或 shard_map 内部使用。为什么？因为它们需要知道通信的「范围」——哪些设备参与通信。

# 一个典型的梯度同步例子
def train_step(params, batch):
    def loss_fn(params):
        logits = model(params, batch['x'])
        return jnp.mean((logits - batch['y']) ** 2)
    
    loss, grads = jax.value_and_grad(loss_fn)(params)
    
    # 跨设备求平均梯度
    grads = jax.lax.pmean(grads, axis_name='batch')
    
    # 更新参数
    params = jax.tree_map(lambda p, g: p - lr * g, params, grads)
    return params, loss

# 用 pmap 包装
parallel_step = jax.pmap(train_step, axis_name='batch')

注意 axis_name='batch' 这个参数。它告诉 JAX：所有在 'batch' 这个轴上切分的设备，都要参与 pmean 通信。这个命名必须和 pmap 或 shard_map 里的轴名一致。

知识体系总览

下面这张图，是我自己梳理的 JAX 分布式核心概念关系。你看一眼，应该能对整体有个把握。

从这张图你可以看到：pmap 和 shard_map 是入口，它们决定了数据怎么分布。数据分布后得到 ShardedDeviceArray。而集体通信原语，则是让这些分散的数据能够协同工作的「桥梁」。

小结

这一章我们聊了 JAX 分布式的三个核心概念：

• pmap：简单粗暴，适合数据并行入门
• shard_map：灵活精细，适合模型并行和混合并行
• ShardedDeviceArray：你看不见的数据分布，但性能影响很大
• 集体通信原语：设备间通信的瑞士军刀

我个人建议，新项目直接上 shard_map。虽然学习曲线陡一点，但后面做模型并行、流水线并行时，你会发现之前的投入都是值得的。

下一章，我们会把这些概念串起来，写一个真正的分布式训练循环。到时候你会看到，这些原语是怎么在实战中配合的。