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 的分布式训练最常用的就是 DDPDistributedDataParallel。它的核心逻辑非常直白每个 GPU 上跑一份完整的模型副本每张卡吃不同的数据反向传播算完梯度之后通过 NCCL 做一次 all-reduce把梯度在所有卡之间平均一下然后每张卡用同样的梯度更新自己那份参数。因为初始参数一样、更新梯度一样所以更新完之后每张卡的参数还是保持一致的。这个思路的好处是心智负担极低。你不需要关心参数怎么切、通信怎么调度只要把模型包一层DDP把 dataloader 换成DistributedSampler基本上就能跑起来。通信和计算的重叠、bucket 的划分、梯度压缩这些优化PyTorch 都在底层帮你做了。但代价也很明显每张卡都存了一份完整的模型参数、优化器状态和梯度。对于一个 7B 参数的模型光参数就是 14GBfp16加上优化器状态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 这边最基础的分布式原语是pmapimport 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, opdist.ReduceOp.SUM)JAX 里对应的操作是jax.lax.psumfrom jax import lax result lax.psum(x, axis_namedevices)但psum必须在pmap的上下文里使用而且需要给pmap指定axis_namejax.pmap(axis_namedevices) def train_step(params, batch): grads jax.grad(loss_fn)(params, batch) grads lax.psum(grads, axis_namedevices) return grads注意这里psum是对整个梯度树做的而 PyTorch 的 all-reduce 是对每个 tensor 单独做的。JAX 的lax.psum支持 pytree所以你可以直接传一个嵌套的字典或列表进去它会自动递归处理。还有一个差异PyTorch 的 all-reduce 是原地操作in-place会修改传入的 tensorJAX 的psum是函数式的返回一个新的值。这个差异看起来小但在写代码的时候影响很大——你不能在 JAX 里写lax.psum(grads, ...)然后指望grads被修改必须写成grads lax.psum(grads, ...)。3.3 从 DistributedSampler 到 shard数据切分的方式PyTorch 里切数据靠DistributedSamplersampler DistributedSampler(dataset, num_replicasworld_size, rankrank) dataloader DataLoader(dataset, samplersampler, batch_sizeper_gpu_batch)JAX 里没有 sampler 这个概念数据切分是靠pmap自动完成的。你只需要把全局 batch 准备好然后pmap会自动把第 0 维切到各个设备上。比如全局 batch size 是 2568 张卡那每张卡实际拿到 32 条数据。但这里有个坑pmap切分的是数组的第 0 维所以你的数据必须已经是一个 shape 为(256, ...)的数组。如果你用的是 Python 的 list 或者 generator需要先转换成jnp.array。而且这个转换是在 host 上做的如果数据量很大可能会成为瓶颈。实践中通常会用tf.data或者torch.utils.data做数据加载然后用jax.pmap的in_axes来指定切分方式。3.4 关键概念对照表概念PyTorchJAX并行原语DDP/FSDPpmap/shard_map通信原语dist.all_reducelax.psum数据切分DistributedSamplerpmap自动切第 0 维参数切分FSDP 自动切PartitionSpec显式声明通信调度运行时动态编译时静态梯度同步DDP 自动手动psum或grad的psum设备网格无显式概念Meshaxis_names这张表里最值得注意的就是最后一行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(), lr1e-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_namejax.pmap(axis_namedevices) def train_step(params, x, y): grads grad(loss_fn)(params, x, y) grads jax.lax.psum(grads, axis_namedevices) 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 层的 MLPhidden size 1024batch size 全局 4096。PyTorch DDP 和 JAX pmap 的吞吐量对比如下框架吞吐量samples/sec显存占用GB/卡PyTorch DDP12,80018.2JAX pmap14,20016.8JAX pmap jit15,60016.8JAX 的吞吐量比 PyTorch 高大概 10% 到 20%显存占用也略低。这个差异主要来自 XLA 的算子融合和通信调度优化。但要注意这个优势在模型结构简单的时候比较明显如果模型里有大量动态控制流或者自定义算子JAX 的优势可能会缩小甚至消失。5. 那些文档里不会写的踩坑记录5.1 pmap 的 axis_name 必须和 psum 一致这是最容易犯的错误之一。pmap的axis_name和psum的axis_name必须完全一致否则会报NameError或者更隐晦的 shape 错误。我见过有人写pmap(axis_namebatch)然后psum(..., axis_namedevices)结果报了一个完全看不懂的错排查了半天才发现是名字对不上。建议统一用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_namedevices)而不是def loss_fn_with_psum(params, x, y): loss loss_fn(params, x, y) return jax.lax.psum(loss, axis_namedevices) 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_namedevices, in_axes(None, 0, 0)) def train_step(params, x, y): ...这里None表示参数不切分因为已经 broadcast 过了0表示 x 和 y 在第 0 维上切分。5.4 常见报错速查表报错信息可能原因解决方法pmap requires axis_name用了psum但pmap没指定axis_name给pmap加上axis_namedevicesaxis_name mismatchpmap和psum的axis_name不一致统一命名shape mismatch in pmap输入的第 0 维不等于设备数检查 batch size 是否能被设备数整除CUDA_ERROR_OUT_OF_MEMORY显存不够减小 batch size 或使用梯度累积NaN in gradients参数初始化不一致或学习率过大检查参数 broadcast 是否正确jax.devices()只返回 CPUGPU 版 JAX 没装对重新安装jax[cuda12]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, meshmesh, in_specs(P(devices), P(devices)), out_specsP(devices)) def train_step(params, x): ...这个代码看起来比pmap复杂但它能表达更灵活的切分策略。比如你可以把参数切到(devices, None)数据切到(None, devices)实现数据并行和模型并行的混合。6. 迁移决策什么时候该用 JAX什么时候留在 PyTorch6.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的位置放错了。