PyTorch DDP 与 JAX pmap 分布式训练差异及迁移实践

1. 从一次多卡训练翻车说起

上周帮一个朋友调一个多卡训练的任务，他之前一直用 PyTorch，这次因为项目需要换成了 JAX。代码写完之后跑单卡没问题，一上 8 卡就各种奇怪的报错——有的是 shape 对不上，有的是梯度莫名其妙变成 NaN，还有一次是 loss 直接卡住不动了。他问我：“不就是换个框架吗，怎么分布式这块差别这么大？”

这个问题其实特别典型。JAX 和 PyTorch 在分布式训练上的差异，不是 API 名字换一换那么简单，而是两套完全不同的设计哲学在打架。PyTorch 的 DDP 是“我告诉你哪些参数要同步，你帮我同步”；JAX 的pmap/shard_map是“我告诉你数据怎么切、参数怎么放，剩下的编译器来搞定”。前者是命令式的、显式的，后者是函数式的、声明式的。

这篇文章就是想把这件事讲透。如果你正在从 PyTorch 迁移到 JAX，或者单纯想搞清楚这两个框架在分布式场景下到底哪里不一样，那接下来的内容应该能帮你少踩几个坑。我会从设计思路、核心 API、实操流程、常见报错几个角度展开，尽量用我实际踩过的例子来说明，而不是照本宣科念文档。

2. 两套分布式哲学的根本分歧

2.1 PyTorch DDP：把同步这件事做到极致简单

PyTorch 的分布式训练，最常用的就是 DDP（DistributedDataParallel）。它的核心逻辑非常直白：每个 GPU 上跑一份完整的模型副本，每张卡吃不同的数据，反向传播算完梯度之后，通过 NCCL 做一次 all-reduce，把梯度在所有卡之间平均一下，然后每张卡用同样的梯度更新自己那份参数。因为初始参数一样、更新梯度一样，所以更新完之后每张卡的参数还是保持一致的。

这个思路的好处是心智负担极低。你不需要关心参数怎么切、通信怎么调度，只要把模型包一层DDP，把 dataloader 换成DistributedSampler，基本上就能跑起来。通信和计算的重叠、bucket 的划分、梯度压缩这些优化，PyTorch 都在底层帮你做了。

但代价也很明显：每张卡都存了一份完整的模型参数、优化器状态和梯度。对于一个 7B 参数的模型，光参数就是 14GB（fp16），加上优化器状态（Adam 的 m 和 v）就是 42GB，再加梯度又是 14GB，单卡 80GB 的 A100 也就勉强放得下。模型再大一点，DDP 就撑不住了，必须上 FSDP 或者张量并行。

2.2 JAX pmap：把“切分”变成一等公民

JAX 的分布式思路从根上就不一样。它的核心抽象是“sharding”——你显式地告诉编译器，这个数组的哪个维度要切到哪些设备上。比如一个 shape 为(8, 1024, 4096)的参数，你可以声明它在第 0 维上切到 8 张卡，那每张卡实际只存(1, 1024, 4096)。

这个声明是通过jax.sharding模块里的Mesh和PartitionSpec来完成的。Mesh定义了一个设备网格，比如(8,)表示 8 张卡排成一维；PartitionSpec则描述每个张量在这个网格上怎么切。编译器拿到这些信息之后，会自动推导出所有中间结果的切分方式，并在必要的时候插入通信操作。

这就带来一个很关键的差异：在 JAX 里，通信不是“额外加的一步”，而是编译过程的一部分。编译器知道整个计算图，所以它可以把通信和计算重叠起来，甚至在某些情况下把 all-reduce 拆解成更高效的 reduce-scatter + all-gather。这是 PyTorch 的 eager 模式做不到的，因为 eager 模式下每一步都是独立执行的，编译器没有全局视野。

2.3 一个类比：手动挡 vs 自动挡

如果打个比方，PyTorch DDP 像是自动挡的车——你踩油门就走，变速箱什么时候换挡你不用管，但你想精细控制换挡时机也做不到。JAX 的 sharding 更像是手动挡——你得自己决定什么时候换挡、换几挡，但一旦你掌握了，就能把车开得又快又省油。

这个类比还能延伸一点：自动挡的车，你换一辆不同型号的，操作方式基本一样；手动挡的车，每辆车的离合行程、挡位手感都不一样，需要适应。JAX 的 sharding 也是这样，不同的 mesh 形状、不同的 partition spec，性能可能差很多，需要你根据具体的模型和硬件去调。

3. 核心 API 对照与关键概念拆解

3.1 从 DDP 到 pmap：入口函数的差异

PyTorch 里启动分布式训练，典型代码是这样的：

import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP dist.init_process_group("nccl") local_rank = int(os.environ["LOCAL_RANK"]) torch.cuda.set_device(local_rank) model = DDP(model.to(local_rank), device_ids=[local_rank])

然后在训练循环里，每张卡吃不同的数据，loss.backward()之后 DDP 会自动做梯度同步。

JAX 这边，最基础的分布式原语是pmap：

import jax import jax.numpy as jnp @jax.pmap def train_step(params, batch): grads = jax.grad(loss_fn)(params, batch) return jax.tree.map(lambda p, g: p - lr * g, params, grads)

pmap的作用是把一个函数编译成可以在多个设备上并行执行的版本。你调用train_step(params, batch)的时候，如果params和batch的第 0 维大小等于设备数，JAX 就会自动把第 0 维切到各个设备上，每个设备执行一份，然后返回的结果也是在第 0 维上堆叠的。

这里有个关键点：pmap要求所有输入和输出的第 0 维都必须等于设备数。这意味着你的 batch 必须能被设备数整除，参数也必须有一个维度能被设备数整除。如果模型参数没有这样的维度，你就得用jax.pmap的in_axes参数来指定哪些参数不切分（用None表示），然后手动做 all-reduce。

3.2 从 all-reduce 到 psum：通信原语的差异

PyTorch 里做 all-reduce 很简单：

dist.all_reduce(tensor, op=dist.ReduceOp.SUM)

JAX 里对应的操作是jax.lax.psum：

from jax import lax result = lax.psum(x, axis_name='devices')

但psum必须在pmap的上下文里使用，而且需要给pmap指定axis_name：

@jax.pmap(axis_name='devices') def train_step(params, batch): grads = jax.grad(loss_fn)(params, batch) grads = lax.psum(grads, axis_name='devices') return grads

注意这里psum是对整个梯度树做的，而 PyTorch 的 all-reduce 是对每个 tensor 单独做的。JAX 的lax.psum支持 pytree，所以你可以直接传一个嵌套的字典或列表进去，它会自动递归处理。

还有一个差异：PyTorch 的 all-reduce 是原地操作（in-place），会修改传入的 tensor；JAX 的psum是函数式的，返回一个新的值。这个差异看起来小，但在写代码的时候影响很大——你不能在 JAX 里写lax.psum(grads, ...)然后指望grads被修改，必须写成grads = lax.psum(grads, ...)。

3.3 从 DistributedSampler 到 shard：数据切分的方式

PyTorch 里切数据靠DistributedSampler：

sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank) dataloader = DataLoader(dataset, sampler=sampler, batch_size=per_gpu_batch)

JAX 里没有 sampler 这个概念，数据切分是靠pmap自动完成的。你只需要把全局 batch 准备好，然后pmap会自动把第 0 维切到各个设备上。比如全局 batch size 是 256，8 张卡，那每张卡实际拿到 32 条数据。

但这里有个坑：pmap切分的是数组的第 0 维，所以你的数据必须已经是一个 shape 为(256, ...)的数组。如果你用的是 Python 的 list 或者 generator，需要先转换成jnp.array。而且这个转换是在 host 上做的，如果数据量很大，可能会成为瓶颈。实践中通常会用tf.data或者torch.utils.data做数据加载，然后用jax.pmap的in_axes来指定切分方式。

3.4 关键概念对照表

这张表里最值得注意的就是最后一行：PyTorch 没有“设备网格”这个显式概念，你通过rank和world_size来隐式地描述设备拓扑；JAX 则要求你显式地定义一个Mesh，并且给每个维度起名字。这个差异在写复杂并行策略（比如同时做数据并行和模型并行）的时候会非常明显。

4. 实操：把 DDP 训练改成 JAX pmap 训练

4.1 环境准备与依赖安装

先说环境。JAX 的安装和 PyTorch 不太一样，它分 CPU 版和 GPU 版，而且 GPU 版需要匹配 CUDA 版本。截至我写这篇文章的时候，JAX 对 CUDA 12 的支持已经比较成熟了，安装命令大概是：

pip install --upgrade "jax[cuda12]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

如果你用的是比较新的显卡（比如 40 系或者 50 系），建议先确认一下 JAX 版本和 CUDA 驱动的兼容性。我遇到过 5060Ti 上装完 JAX 之后jax.devices()只返回 CPU 的情况，最后发现是 CUDA 版本不匹配，换成jax[cuda12_pip]才解决。

验证安装是否成功：

import jax print(jax.devices())

如果输出里有CudaDevice并且数量等于你的 GPU 数，那就没问题。如果只有CpuDevice，说明 GPU 版没装对。

4.2 单卡训练代码回顾

先看一个 PyTorch 的单卡训练代码，这是我们后面要改写的基准：

import torch import torch.nn as nn model = nn.Linear(1024, 1024).cuda() optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) for batch in dataloader: x, y = batch x, y = x.cuda(), y.cuda() optimizer.zero_grad() loss = nn.functional.mse_loss(model(x), y) loss.backward() optimizer.step()

对应的 JAX 单卡版本：

import jax import jax.numpy as jnp from jax import grad, jit def loss_fn(params, x, y): pred = jnp.dot(x, params['w']) + params['b'] return jnp.mean((pred - y) ** 2) @jit def train_step(params, x, y): grads = grad(loss_fn)(params, x, y) return jax.tree.map(lambda p, g: p - lr * g, params, grads) params = {'w': jnp.ones((1024, 1024)), 'b': jnp.zeros(1024)} for x, y in dataloader: params = train_step(params, x, y)

注意 JAX 里没有optimizer.zero_grad()这一步，因为梯度是函数式返回的，不存在“累积”的概念。也没有loss.backward()，因为grad直接返回梯度。这是函数式编程和命令式编程的一个典型差异。

4.3 改写成 pmap 版本

现在把上面的 JAX 单卡代码改成pmap版本。关键改动有三处：

第一，给train_step加上pmap装饰器，并指定axis_name：

@jax.pmap(axis_name='devices') def train_step(params, x, y): grads = grad(loss_fn)(params, x, y) grads = jax.lax.psum(grads, axis_name='devices') grads = jax.tree.map(lambda g: g / jax.device_count(), grads) return jax.tree.map(lambda p, g: p - lr * g, params, grads)

第二，把参数复制到所有设备上。pmap要求输入的第 0 维等于设备数，所以我们需要把参数 broadcast 一下：

params = jax.tree.map(lambda p: jnp.broadcast_to(p, (jax.device_count(),) + p.shape), params)

第三，把数据也切成(device_count, per_device_batch, ...)的形状：

x = x.reshape(jax.device_count(), -1, *x.shape[1:]) y = y.reshape(jax.device_count(), -1, *y.shape[1:])

然后调用train_step(params, x, y)就可以了。返回的params也是在第 0 维上堆叠的，取params[0]就是第一张卡上的参数。

这里有个细节：psum之后我加了一步除以device_count，这是因为psum是求和，而 DDP 的 all-reduce 默认是求平均。如果你希望和 DDP 行为一致，就需要手动除一下。当然你也可以用jax.lax.pmean，它直接就是求平均。

4.4 参数初始化的坑

上面那个 broadcast 参数的写法看起来简单，但实际用的时候有个坑：如果你的参数是通过jax.random生成的，每张卡上的初始值可能不一样。因为jax.random的随机数生成器是函数式的，你需要显式地指定相同的 seed，或者用jax.random.split之后把同一个 key 复制到所有设备上。

我踩过的一个坑是：用jax.random.normal(key, shape)生成参数，然后直接pmap，结果每张卡上的初始参数不一样，训练的时候 loss 一直震荡。后来改成先用单卡生成参数，再 broadcast 到所有设备，问题就解决了。

key = jax.random.PRNGKey(42) params = jax.random.normal(key, (1024, 1024)) params = jnp.broadcast_to(params, (jax.device_count(),) + params.shape)

4.5 性能对比与实测数据

我在 8 张 A100 上做了一个简单的对比实验，模型是一个 4 层的 MLP，hidden size 1024，batch size 全局 4096。PyTorch DDP 和 JAX pmap 的吞吐量对比如下：

JAX 的吞吐量比 PyTorch 高大概 10% 到 20%，显存占用也略低。这个差异主要来自 XLA 的算子融合和通信调度优化。但要注意，这个优势在模型结构简单的时候比较明显，如果模型里有大量动态控制流或者自定义算子，JAX 的优势可能会缩小甚至消失。

5. 那些文档里不会写的踩坑记录

5.1 pmap 的 axis_name 必须和 psum 一致

这是最容易犯的错误之一。pmap的axis_name和psum的axis_name必须完全一致，否则会报NameError或者更隐晦的 shape 错误。我见过有人写pmap(axis_name='batch')然后psum(..., axis_name='devices')，结果报了一个完全看不懂的错，排查了半天才发现是名字对不上。

建议：统一用'devices'或者'i'这种简短的名字，并且在代码里全局搜索确认没有拼写错误。

5.2 psum 的梯度问题

lax.psum本身是可微的，但它的梯度行为和你想的可能不一样。psum的梯度是psum的转置，也就是psum本身（因为psum是线性操作）。这意味着如果你在grad函数内部调用psum，梯度会被重复计算。

正确的做法是在grad之后调用psum，就像我上面写的那样：

grads = grad(loss_fn)(params, x, y) grads = jax.lax.psum(grads, axis_name='devices')

而不是：

def loss_fn_with_psum(params, x, y): loss = loss_fn(params, x, y) return jax.lax.psum(loss, axis_name='devices') grads = grad(loss_fn_with_psum)(params, x, y) # 这样梯度会不对

5.3 数据加载的瓶颈

JAX 的pmap要求数据已经是一个完整的数组，这意味着你必须在 host 上把数据准备好，然后一次性传给设备。如果数据量很大，host 到 device 的传输会成为瓶颈。PyTorch 的DataLoader可以在后台线程里预取数据，JAX 这边需要你自己用tf.data或者torch.utils.data配合jax.pmap的in_axes来实现类似的效果。

我常用的一个模式是用tf.data做数据加载和预处理，然后用jax.pmap的in_axes指定数据切分：

@jax.pmap(axis_name='devices', in_axes=(None, 0, 0)) def train_step(params, x, y): ...

这里None表示参数不切分（因为已经 broadcast 过了），0表示 x 和 y 在第 0 维上切分。

5.4 常见报错速查表

5.5 从 pmap 到 shard_map 的演进

pmap虽然简单，但它的表达能力有限。它只能做数据并行，而且要求所有输入输出的第 0 维都等于设备数。如果你需要做模型并行（比如把一个大矩阵切到多张卡上），pmap就不够用了。

JAX 现在推荐的做法是用shard_map配合jax.sharding.Mesh。shard_map允许你显式地指定每个张量的切分方式，而且支持更复杂的并行策略。不过shard_map的学习曲线比pmap陡峭不少，建议先把pmap用熟了再过渡。

from jax.sharding import Mesh, PartitionSpec as P from jax.experimental.shard_map import shard_map mesh = Mesh(jax.devices(), axis_names=('devices',)) @partial(shard_map, mesh=mesh, in_specs=(P('devices'), P('devices')), out_specs=P('devices')) def train_step(params, x): ...

这个代码看起来比pmap复杂，但它能表达更灵活的切分策略。比如你可以把参数切到('devices', None)，数据切到(None, 'devices')，实现数据并行和模型并行的混合。

6. 迁移决策：什么时候该用 JAX，什么时候留在 PyTorch

6.1 JAX 的优势场景

如果你的模型结构比较规整（比如全是 Transformer 或者 MLP），而且你愿意花时间调 sharding 策略，JAX 通常能带来更好的性能。XLA 的算子融合和通信调度优化在规整的计算图上效果非常明显。另外，JAX 的函数式特性让它在做元学习、强化学习这类需要高阶梯度的场景下更自然。

还有一个容易被忽略的点：JAX 的jit编译之后，整个训练步骤是一个单一的 XLA 计算图，这意味着你可以用 XLA 的 profiling 工具看到每一步的详细耗时，包括通信和计算的重叠情况。PyTorch 的 profiler 虽然也能看，但粒度没有 XLA 那么细。

6.2 PyTorch 的优势场景

如果你的模型里有大量动态控制流（比如if分支、while循环），或者你依赖很多自定义的 CUDA 算子，PyTorch 的 eager 模式会更灵活。JAX 的jit对动态控制流的支持有限，虽然可以用lax.cond和lax.while_loop来写，但代码会变得很别扭。

另外，PyTorch 的生态更成熟，各种预训练模型、工具库、调试工具都比 JAX 丰富。如果你需要快速原型验证，或者团队里大部分人只熟悉 PyTorch，那留在 PyTorch 可能是更务实的选择。

6.3 混合使用的可能性

其实这两个框架不是非此即彼的关系。我见过一些项目用 PyTorch 做数据加载和预处理，然后用 JAX 做模型训练，通过dlpack在两者之间传递张量。这种混合方案在需要利用 JAX 的性能优势但又不想放弃 PyTorch 生态的时候比较有用。

不过这种方案也有代价：数据在 host 和 device 之间来回拷贝的开销可能抵消掉 JAX 的性能优势。而且调试起来更麻烦，因为你要同时理解两个框架的行为。

7. 我个人的一些经验体会

从 PyTorch 迁移到 JAX 的分布式训练，最大的障碍不是 API 不熟悉，而是思维方式的转变。PyTorch 的 DDP 让你觉得“分布式”是一个可以后期加上的东西——你先写好单卡代码，然后包一层 DDP 就完事了。JAX 则要求你从写第一行代码的时候就考虑切分，因为pmap和sharding是渗透到整个计算图里的。

我自己的经验是，如果你决定用 JAX 做分布式训练，最好一开始就用pmap或者shard_map来写，不要先写单卡再改。因为单卡代码里的很多写法（比如用 Python 的for循环遍历参数、用if判断 shape）在pmap下会直接报错，改起来很痛苦。

还有一个很实用的技巧：在调试pmap代码的时候，先用jax.config.update('jax_disable_jit', True)把jit关掉，这样报错信息会清晰很多。等逻辑跑通了再打开jit看性能。这个技巧帮我省了很多排查时间。

最后说一个关于psum的细节：如果你在pmap里用了psum，但忘记在grad之后调用，而是直接在loss上调用，梯度会变成所有设备上梯度的和，而不是平均。这个错误很隐蔽，因为 loss 看起来是正常的，但训练出来的模型效果会差很多。我建议在写完pmap代码之后，先用一个小数据集跑几个 step，对比一下单卡和pmap的 loss 曲线，如果差异很大，大概率就是psum的位置放错了。