这几年分布式训练里JAX 和 PyTorch 的“路线之争”是我和同行聊得最多的话题。很多人一上来就问JAX 分布式训练到底和 PyTorch 有什么不一样我用 8 卡 A100 实际跑过 GPT 级别的训练和对比实验后最大的感受是——差异不只是 API 换了一套而是你思考模型、数据、设备的方式全变了。这篇文章我会从底层同步机制、代码实战、迁移踩坑、性能调优四个层面把这两种框架的分布式训练彻底拆开讲清楚。无论你是想从 PyTorch 迁移到 JAX还是单纯想搞懂为什么 JAX 能写出更简洁的分布式代码这篇文章都值得看完。1. 分布式训练的本质你同步的是什么1.1 PyTorch DDP 的数据并行哲学PyTorch 的分布式训练绝大多数人用的是 DistributedDataParallelDDP。它的核心思路非常直白每个 GPU 上放一份完整的模型副本每个进程读取不同的 batch 数据做前向和反向计算计算出梯度后通过 all-reduce 操作把梯度在所有进程间做求平均然后每个进程用平均后的梯度更新自己的模型副本。听起来很简单但实际工程里有很多细节。DDP 在反向传播过程中会给每个参数的梯度挂一个 hook当某个参数的梯度计算完成后hook 立即触发把梯度做异步 all-reduce。这样梯度的通信和下一层的反向计算重叠起来能隐藏一部分通信延迟。同步方式是 N 个进程之间的集合通信底层用的是 NCCL 库。NCCL 在单机多卡上会优先走 NVLink多机再走 InfiniBand 网络。这里有个关键点DDP 同步的是梯度而不是参数。模型参数在每台机器上是完全冗余的。每次迭代结束后所有进程的梯度达到一致但参数的一致性其实依赖的是“用相同梯度做相同更新”这个隐含假设。只要学习率、优化器状态一致参数就会保持一致。这也是 DDP 在 PyTorch 里长时间是默认方案的原因——它几乎不动用户原有的单机代码只是把你的 model 包一层 model DDP(model)然后保证每个数据加载器喂给不同 GPU 的数据不重合就行。但冗余参数带来的问题也明显显存里每个卡都要放一份完整的模型状态。如果模型规模超过单卡显存DDP 就趴窝了。所以后来社区在 DDP 之上又搞出了 FSDPFully Sharded Data Parallel把参数、梯度、优化器状态都分片到各卡上。FSDP 本质上是把模型参数存储和通信都做了分片这跟 JAX 的思维其实有相通之处。1.2 JAX 的 SPMD 与 XLA 编译JAX 分布式训练的底层逻辑完全不同。它的核心是函数式纯函数 XLA 编译。在 JAX 里你定义一个计算函数比如一个 loss 函数函数输入是参数和 batch 数据输出是 loss 和梯度。这个函数本身不带任何全局状态也不修改任何参数。分布式是用 pmap 这个 API 来实现的pmap 会把同一个函数复制到多个设备上并且在每个设备上喂给它不同的 batch 切片然后自动在全设备之间做必要的通信。pmap 的含义是 Single Program Multiple DataSPMD。被 pmap 包裹的函数会在 N 台设备上并行执行函数的第一个参数通常是一个 batch会被切分成 N 份分发到每一台设备。函数内部如果出现了类似“对所有设备上的数据做全局规约”这种操作需要显式地通过 lax.pmean 这样的集合通信原语去写清楚。如果没有这样的通信那么各设备互不打扰完全并行。这里最关键的是JAX 的分布式是在编译期决定的不是运行时才去同步。你用 jit 编译一个函数时XLA 编译器会分析整个计算图决定每个张量放在哪台设备上、在哪个节点插入通信算子。JAX 里有一个自动分片系统叫 GSPMD它不需要你手写 pmap只要你在 jit 里通过设备布局约束比如用 jax.sharding.NamedSharding 指定某个数组的分片方式编译器就能自动推导出所有中间结果的分片方案自动插入 all-gather、reduce-scatter 这些通信。所以你会发现JAX 的分布式更像“编译器帮你规划数据布局和通信”而 PyTorch 的 DDP 则是“运行时通过 hook 触发同步”。这两种哲学听起来都行但实际写代码时的思维方式完全不同。1.3 思维模型对比显式控制 vs 编译器接管为了把差异说得更清楚我给你画一个直观的对比表格维度PyTorch DDPJAX (pmap / jit GSPMD)同步对象梯度取决于你的计算图可能是梯度、参数、中间激活同步时机反向传播过程中运行时编译期静态分析后插入通信算子数据并行方式每个进程一份完整模型副本模型副本也可以分片数据天然分片代码风格面向对象模型有状态函数式一切皆为纯函数调试方式可以 print 中间结果可以断点需要将函数转移到 CPU 或用 debug 回调扩展能力DDP → FSDP 是逐步演进在 jit 里一句设备约束即可实现多种并行策略我自己的体会是PyTorch 的 DDP 是“手动挡”你需要自己握好方向盘分配数据、踩离合器调用 hook 同步JAX 更像一辆聪明的“自动驾驶”你跟它说目的地定义函数和分片约束它自己规划路径XLA 生成通信计划。但自动驾驶也有问题——它有时候会走一条你没预料到的路线比如自动插入你不知道的通信算子导致性能不符合预期。这就得靠经验去调整分片约束。有一个常见的误解是 JAX 的 pmap 就是像 PyTorch 的 DataParallel。不对。pmap 不要求模型参数必须在每台设备上有一个完整副本你完全可以把不同层参数放到不同设备上然后通过 lax.psum 等循环调用做跨设备规约。这天然支持了模型并行、流水线并行、数据并行混合的姿态。PyTorch 里这些你要么用 FSDP要么自己手写张量并行工作量大不少。2. 代码实战两种框架的分布式训练实现2.1 PyTorch DDP 的标准写法先看 PyTorch 的经典 DDP 训练主循环。我直接给一段能跑的简化代码import os import torch import torch.distributed as dist import torch.nn as nn import torch.optim as optim from torch.nn.parallel import DistributedDataParallel as DDP def setup(): dist.init_process_group(backendnccl) def cleanup(): dist.destroy_process_group() def train(): rank dist.get_rank() world_size dist.get_world_size() torch.cuda.set_device(rank) model nn.Linear(1024, 1024).cuda() model DDP(model, device_ids[rank]) optimizer optim.SGD(model.parameters(), lr0.01) data torch.randn(64, 1024).cuda() # 实际会从DataLoader取 for step in range(100): optimizer.zero_grad() loss model(data).sum() loss.backward() optimizer.step() if rank 0 and step % 10 0: print(fstep {step}, loss {loss.item():.4f}) cleanup() if __name__ __main__: setup() train()启动方式通常是 mp.spawn 或 torchrun --nproc_per_node8 train.py。这段代码里有几个关键点值得注意第一model.cuda() 必须在 DDP 包装之前完成第二DDP 会自动给模型梯度加 hook但你必须调用 loss.backward() 才会触发 all-reduce第三所有进程的执行流程完全相同靠数据加载器保证喂给每个进程的数据不一样。这些细节在初期很容易踩坑最典型的是忘了 set_device导致所有进程都往 GPU0 上写直接显存爆炸。DDP 的 hook 机制有个实际好处反向传播过程中梯度一算完就开始通信通信和计算重叠。如果你在 backward() 之后手动调用 all-reduce那通信和计算就变成了串行性能会掉 20% 以上。所以千万别自作聪明去改 DDP 的梯度同步逻辑。2.2 JAX pmap 的标准写法JAX 的代码风格明显不一样。先看一个用 pmap 做最简单分布式训练的示例import jax import jax.numpy as jnp from jax import random, grad from jax.lax import pmean def loss_fn(params, x, y): pred x params return jnp.mean((pred - y) ** 2) jax.jit jax.pmap def train_step(params, x, y, lr): loss, grads jax.value_and_grad(loss_fn)(params, x, y) grads pmean(grads, axis_namedevices) # 跨设备平均梯度 new_params params - lr * grads return new_params, loss rng random.PRNGKey(0) params random.normal(rng, (1024, 1024)) # 假设有8个设备每个设备喂入 16 个样本 x random.normal(rng, (8, 16, 1024)) y random.normal(rng, (8, 16)) for step in range(100): params, loss train_step(params, x, y, 0.01)理解这段代码要知道三个概念value_and_grad 同时算出损失和梯度纯函数式没有任何累加状态pmean 是跨设备做平均axis_name 必须是 pmap 里指定的设备轴名pmap 会自动把你传入的 x 的第一维8切成 8 份每个设备拿一份。注意 train_step 被 jit 和 pmap 同时装饰的顺序。装饰顺序从下往上执行所以先 pmap 复制到设备再整体编译。这个函数一次编译之后后续每 step 调用都直接执行编译后的机器码非常快。但这里有一个容易混淆的地方pmap 里的“第一维”必须适配设备数。比如你有 8 张卡x 的第一维必须是 8 或者 8 的倍数。如果写成 16jax 会直接报错告诉你 batch 维度必须等于设备数。这就是 JAX 跟 PyTorch 在代码自由度上的差异——PyTorch 会偷偷帮你在每个进程取不同数据JAX 则是强迫你把 batch 显式地按设备切好。2.3 进阶jax.shard_map 与 FSDP 的对应关系pmap 在模型很大时也会遇到显存问题。因为它默认在每台设备上也保留一份完整参数副本。想做得更精细必须手动分片。JAX 后面推出了 jax.shard_map它允许你显式指定每个张量在设备网格上的分片布局更贴合 FSDP 的思维。举个例子如果我有 8 台设备想做一个 2 维设备网格每维 4 台然后把权重矩阵按行分片到 grid[0] 上按列分片到 grid[1] 上我可以这样写from jax.sharding import Mesh, PartitionSpec, NamedSharding import jax.numpy as jnp devices jax.devices() # 8台设备 mesh Mesh(devices.reshape(2, 4), (data, model)) sharding NamedSharding(mesh, PartitionSpec(data, model)) x jax.device_put(jnp.zeros((8, 16)), sharding)这样 x 会被切分到整个设备网格上。后续用 jit 编译时XLA 的 GSPMD 会自动推断出所有中间张量的分片并插入必要的通信。这种自动分片的自由度远超 PyTorch FSDP因为 FSDP 还是以“每个模型层”为粒度去分片而 JAX 可以精确到每个张量的每个维度。用 shard_map 还可以实现“手动并行策略”你可以把 transformer 的 self-attention 里的 q/k/v 分到不同设备组把 FFN 分到另一组设备组。这在 PyTorch 里通常要用 Megatron 那样的框架但 JAX 里你只写几个 PartitionSpec 就行。当然代价是你要理解分片规格的含义否则很容易写出通信量爆炸的配置。3. 从 PyTorch 迁移到 JAX 的实操记录3.1 环境搭建的坑我最早是在一台 8 卡 A100 上做迁移测试的。先说说安装。你可能已经熟悉了“安装pytorch”时用 conda 激活环境、再 pip install torch 那一套。JAX 的安装跟 PyTorch 完全不同它官方推荐直接在虚拟环境里用 pip 装带有 CUDA 支持的版本pip install -U jax[cuda12]注意这里不需要你自己折腾 CUDA toolkitpip 包会把 JAX 依赖的 CUDA 运行时库一并带上。这意味着 JAX 的安装对 CUDA 驱动有一定的“黑盒”依赖如果驱动版本太老很容易出现 CUDA_ERROR_NO_DEVICE。所以我的建议是先跑一下 nvidia-smi确认驱动支持 CUDA 12.x再装 jax[cuda12]。如果用老驱动就装 jax[cuda11]。另外JAX 默认只在 GPU 上跑单个设备你要多卡分布式不需要额外装什么 Horovod、DeepSpeed但需要确保 jax.devices() 能识别出所有 GPU。如果识别不到检查一下 JAX 是用哪个 CUDA 版本编译的再对比驱动支持的版本。这里和 PyTorch 安装的排查思路差不多但 PyTorch 的报错更直接JAX 经常是“静默地只用 CPU”让你误以为程序跑错了。我把环境搭建的坑列在下面全是实测过的现象原因解决jax.devices() 返回 CPU 设备装了 CPU 版 JAX或 CUDA 库不匹配pip install -U jax[cuda12] 并检查驱动多卡设备不识别没有安装 NVIDIA 容器工具包或驱动限制确认 nvidia-smi 能看到全部 GPU重启服务OSError: libcudnn.so.8 不存在JAX 包带的 cuDNN 与系统库冲突用 pip install -U jaxlib 升级或切换版本编译超级慢第一次 jit 需要编译 XLA 图正常现象编译后缓存到 /tmp重启后重新编译3.2 核心代码迁移从 DDP 到 pmap 的思维转变我迁移的第一个模型是一个简单的 MLP。用 PyTorch 写非常顺模型定义成一个类参数藏在 nn.Module 里训练循环里随时更新状态。但 JAX 里必须重新定义参数是字典或 pytree嵌套的数组容器前向函数是纯函数优化器状态要自己显式持有。从 PyTorch 到 JAX 的迁移我最崩溃的就是下面这些差异在 PyTorch 里你可以放心地做 model.eval() 来切换 BN 层行为但在 JAX 里如果你有 BatchNorm它的 running_mean 是状态必须作为参数显式传进函数再传出来。每个训练 step 返回新状态然后通过更新参数来延续。这非常像函数式语言里用“不可变状态”实现循环。还有一点PyTorch 的优化器比如 Adam内部有 momentum 和 variance这些状态在 JAX 里也需要你自己维护。还好 optax 这个库帮你封装好了用法类似optax.adam(learning_rate)然后通过opt.init(params)创建优化器状态opt.update返回更新后的状态和梯度。初始迁移时我差点自己手写 Adam后来发现用 optax 就顺了import optax optimizer optax.adam(0.001) opt_state optimizer.init(params) jax.jit jax.pmap def train_step(params, opt_state, x, y): loss, grads jax.value_and_grad(loss_fn)(params, x, y) grads pmean(grads, axis_namedevices) updates, new_opt_state optimizer.update(grads, opt_state, params) new_params optax.apply_updates(params, updates) return new_params, new_opt_state, loss每次 step 后你要用返回的 new_params 替换旧 params。这种显式传递状态的写法对 PyTorch 用户来说一开始非常别扭但坚持两三天后你会发现它有一个巨大的优势整个训练过程没有任何隐式状态你随时可以把中间参数保存下来也可以把函数拿到 CPU 上单步调试甚至可以把一次训练恢复成“无状态的执行轨迹”。3.3 实测性能数据与调优思路我拿同样的 MLP1024 神经元4 层分别用 PyTorch DDP 和 JAX pmap 在 8 卡 A100 上跑 1000 步batch 大小一样。结果很有意思JAX 的第一步跑了 40 秒编译 XLA之后每步 0.8 毫秒PyTorch DDP 第一步很快开始了但稳定后每步 1.2 毫秒。也就是说JAX 在小模型上有性能优势但编译开销需要你忍一次。如果模型大一点比如像 GPT-2 那样的规模JAX 的优势会更明显因为 XLA 可以融合更多算子通信计划也更紧凑。但要注意JAX 的 pmap 默认在每步同步梯度如果你用的是 pmean它等价于 DDP 的 all-reduce。但如果你想做梯度累积PyTorch 里直接调 loss.backward() 把梯度累加到 .grad 里而 JAX 里你不能跨步累积梯度——因为纯函数不允许修改全局状态。要实现梯度累积你只能在函数内部用一个循环迭代累积再把累积梯度应用一次更新。代码会变成def accumulate_step(params, opt_state, x, y, grad_accum): # 在函数内部用 lax.scan 或 for 循环累积梯度 pass看起来复杂但好处是 XLA 可以把你积累的操作完全展开并优化通信量比 PyTorch 的多次 all-reduce 更小。还有一点实际调优经验JAX 里不要轻易用jax.jit包一层巨大的函数如果函数里有大量 Python 循环XLA 的编译时间会急剧上升。我要么把 Python 循环改成 jax.lax.scan 或者 jax.lax.fori_loop要么把它分成多个小 jit。这个取舍在写大规模模型时非常关键。4. 常见问题、踩坑记录与选型建议4.1 问题速查表这里把我这两个月被问得最多、自己也踩过的坑整理成一张速查表基本覆盖了“JAX 分布式训练”从入门到跑通的常见问题。问题原因解决办法pmap 报错 batch 维度不等于设备数pmap 默认把第一个维度按设备数切分把数据 reshape 成 (设备数, 每设备batch, ...)pmean 报错 axis_name 未定义忘了在 pmap 装饰器里指定 axis_name 参数jax.pmap(axis_namedevices)训练 loss 不下降params、grads 没做跨设备平均确认 grad 后调 pmean且 pmean 必须在 pmap 内部编译时间占了一大半函数计算图过于复杂减少 jit 函数里的 Python 循环多用 lax.scan显存 OOMpmap 默认每设备保留完整参数副本使用 jax.shard_map 或 NamedSharding 做参数分片模型保存后加载不回来没保存 opt_stateJAX 要保存 params 和 opt_state两者缺一不可多机通信比单机还慢未配置 RDMA / 网络带宽问题检查 NCCL 环境变量或改用 JAX 的 multihost 配置莫名其妙的 NaN loss没有做梯度裁剪且学习率太大JAX 用 optax.clip_by_global_norm 实现全局梯度裁剪这里面有两个坑特别隐蔽。第一个是 pmean 必须放在 pmap 内部如果你把 pmean 放在 pmap 外面它会把所有设备数据规约到一个标量上然后广播回来结果就是你得到的梯度不是按 batch 平均而是按全局所有样本平均学习率瞬间变大 8 倍整个训练直接崩溃。第二个坑是 JAX 里所有随机数必须显式管理。你在 PyTorch 里用 torch.randn 随便生成数据但在 JAX 里如果不手动传入 PRNGKey那么每个设备生成的随机数可能一致导致数据完全没有随机性loss 卡在同一个值上不动。4.2 什么情况选 JAX什么情况选 PyTorch问得最多的一句话是“那我到底该学哪个”我给一个特别务实的回答看你要解决什么问题。如果你还在深度学习入门阶段或者主要工作是做业务模型部署、快速复现论文那先用好 PyTorch 是几乎必然的选择。PyTorch 的生态太成熟了任何一个 tricks 几乎都有现成实现。分布式场景里DDP 配 FSDP 已经能覆盖 90% 的训练需求社区的踩坑文档也远比 JAX 丰富。但如果你是在做算法研究尤其是需要频繁自定义分布式策略或者你要搞大规模预训练模型JAX 值得你认真投入。XLA 的自动分片和函数式纯函数让你能很清晰地描述“参数放哪、数据怎么切、通信怎么走”这种细粒度控制在 PyTorch 里得写很多底层代码。实际上 Google 的很多新模型都用 JAX 做训练比如一些开源的大模型项目已经直接用 jax shard_map 重构了。另外看团队情况。如果团队里都是 PyTorch 背景强行推 JAX 的学习曲线非常陡峭。我见过好几个团队因为迁移 JAX 后写不出正确的调试代码又退回 PyTorch。这并不丢人工具永远为业务服务。JAX 最大的痛点是“一旦运行报错堆栈信息跟 Python 几乎没关系”你会面对一堆 XLA 内部抽象排查效率很低。这点 PyTorch 的 Python 堆栈友好太多。我的建议是你可以把 JAX 当作理解分布式训练底层原理的“教科书”把 shard_map 和 GSPMD 的概念搞懂之后再回去看 PyTorch 的 FSDP很多设计一下子就能看明白。反过来你有 PyTorch 经验后再接触 JAX也会发现它的 SPMD 抽象没那么神秘。最后分享一个我目前在用的小技巧在 JAX 里调试分布式逻辑不要一上来就跑 8 卡。先用jax.local_devices()里只挑出 2 个设备调试分片逻辑再用jax.debug打印中间张量的sharding属性。一旦你能在 2 卡上跑通8 卡和 32 卡基本就顺了。相比之下PyTorch 用 1 卡到 2 卡调试 DDP 时还得处理进程通信的复杂性JAX 这种“编译期静态检查”反而让你更早点发现问题。踩过几次坑之后我现在的习惯是先用一个小模型把分布式规格验证对再放大模型和 batch——这个习惯让我少熬了不少夜。
