上个月把公司内部一个多卡训练任务从 PyTorch 迁移到了 JAX起因倒不是跟风而是模型里频繁涉及自定义梯度、大批量数据并行以及尝试在单机多卡上跑一些张量并行的实验。折腾了大概两周踩了不少坑也把 JAX 分布式训练和 PyTorch 的底层差异摸得比较清楚了。这篇文章不打算写成翻译文档我想直接聊聊两者在分布式的设计哲学上有多不一样以及在实操中这些差异是怎么体现出来的。无论你是正在观望要不要学 JAX还是已经在 PyTorch 里用 DDP 跑惯了多卡这篇文章都值得你花十分钟过一遍——它至少能帮你想明白一个问题换框架到底换的是 API 还是思维方式1. JAX 和 PyTorch 在底层哲学上的分岔路1.1 JAX 本质上是“可微分的 NumPy”不是“另一个 PyTorch”很多人第一次打开 JAX 文档会下意识拿它和 PyTorch 对标然后发现它连nn.Module都没有第一反应是“这玩意儿能写模型吗”。其实 JAX 对自己的定位从来不是一个轻量级深度学习框架而是一套基于 XLA 的数组计算与自动微分系统它的核心组件是jnpNumPy API、grad自动微分、jit即时编译和vmap向量化映射。模型只是在这些基础组件之上叠加出来的“用法”而不是框架内置的抽象。这个定位差异直接影响了分布式训练的形态。PyTorch 的思维是“命令式 模块化”模型是一个有状态的对象分布式训练里的通信、梯度同步、参数广播都是围绕这些对象做的。JAX 的思维是“函数式 变换”模型只是一个纯函数params - output分布式并行本身也是函数变换的一种。所以你看 JAX 文档里几乎不会出现“训练循环”这种词它一直在强调的是jit、vmap、pmap、pjit这些变换怎么组合。这个区别说起来抽象落到代码上就非常直观。PyTorch 里你要想给某个模块做参数平均得显式写出model.parameters()遍历、all_reduce之后除以世界大小JAX 里你只需要在pmap里写p.mean(params, axis_namebatch)通信自动产生。因为 JAX 把“跨设备并行”这件事变成了函数语义的一部分编译器XLA在编译时就能看到这个并行边界从而做全局优化而不是像 PyTorch 那样在运行时用通信原语把所有设备拼起来。1.2 不可变数组和纯函数给分布式带来了什么JAX 一切计算都建立在不可变数组上——jnp.ndarray没有 in-place 操作所有“更新”都是复制返回新数组。第一次接触会觉得限制很大但放在分布式场景里这个限制其实是极大的优势。因为不可变性意味着没有隐式共享状态一个张量在设备 A 上是什么样它在设备 B 上就是什么样不会出现某一步算完被悄悄改掉的情况。纯函数带来的好处更明显。PyTorch 的 DistributedDataParallel 之所以要 hook 前向和反向是因为模型参数是可变对象必须在运行时去拦截梯度传播、做all_reduce。JAX 里训练步就是一个函数输入params、batch、rng_key输出新的params和 loss。因为这个函数不访问任何外部状态XLA 可以把整步计算压成一个编译好的 kernel 图然后投放到多卡上同步执行。通信模式不再是一堆分散的all_reduce调用而是编译图里预先规划好的集合通信流程。实际跑起来之后你会发现JAX 的多卡训练通常不存在“有一张卡落后然后大家都在等它”的严重情况。PyTorch 里这种 straggler 问题很常见因为你没法预知各卡的 kernel 启动顺序和通信延迟。JAX 因为是 AOT 编译所有通信都在图上XLA 能自动调整各卡的计算排布来减小等待。当然这不等于 JAX 没有性能问题它把问题挪到了编译时间和显存管理上后面我会专门讲。2. 分布式训练的核心差异DDP 的命令式同步 vs JAX 的函数式并行2.1 PyTorch DDP 的 ring-allreduce 到底做了什么PyTorch 里多卡训练的地基是 DistributedDataParallel底层依赖 NCCL 的 ring-allreduce。每张卡各自算梯度然后在进程组里做梯度归约ReduceScatter AllGather最后每个 rank 都拿到全量梯度并更新参数。这个机制很成熟你只需要在模型外面包一层DDP(model)然后照常写loss.backward()和optimizer.step()剩下的事框架都帮你做了。但“照常写”背后有一个隐性成本通信模式是在运行时动态决定的。DDP 并不知道你的模型有多少参数、梯度怎么流动、卡间拓扑长什么样它只是在反向传播结束时无脑做一次全量梯度同步。多机场景下你还得自己关心等级划分rank、local_rank、init_process_group、dist.barrier()这类进程管理细节。通信和计算是两套体系计算是 Python 驱动的 kernel 序列通信是异步的原语调用两者之间靠 hooks 和事件流去对齐。我早年用 DDP 跑过 64 卡的模型并行实验最难受的就是显式切分参数、自己算 rank 对应的切片位置、手动拼all_reduce和all_gather。一旦模型结构变了这些通信代码全部要重写。后来用 FSDP 稍微好点但 FSDP 的分片策略和重新聚合逻辑也是显式配置的仍然需要理解sharding_factor之类的概念。2.2 JAX 的 pmap 把“并行”包装成了函数变换JAX 的pmap全称是 parallel map。你写一个处理单个 batch 的纯函数pmap会自动把它复制到所有设备上然后对 batch 维度做切分函数内部若用了lax.pmean之类的集体操作只要声明axis_nameXLA 就会自动插入对应的通信原语。所以从用户视角看分布式训练和多设备之间的通信几乎消失在了函数语义里。来看一个直观对比。同样是 8 卡数据并行训练PyTorch 的启动方式通常是# PyTorch DDP import torch.distributed as dist dist.init_process_group(backendnccl) local_rank int(os.environ[LOCAL_RANK]) torch.cuda.set_device(local_rank) model MyModel().to(local_rank) model DDP(model, device_ids[local_rank], output_devicelocal_rank) for epoch in range(epochs): for x, y in train_loader: x, y x.to(local_rank), y.to(local_rank) optimizer.zero_grad() loss model(x, y) loss.backward() optimizer.step()这里你虽然不需要管梯度同步但要管进程组、设备映射、数据加载的分布。JAX 的 pmap 版更接近“只写核心计算”# JAX pmap import jax from jax import lax pmap def train_step(params, x, y, key): def loss_fn(p): pred model_apply(p, x) return jnp.mean((pred - y) ** 2) grads grad(loss_fn)(params) grads lax.pmean(grads, axis_namebatch) # 自动跨卡归约 new_params tree_map(lambda p, g: p - lr * g, params, grads) return new_params for epoch in range(epochs): # x 形状是 (num_devices, per_device_batch, ...) # params 已经在各卡上各有一份 params train_step(params, x, y, rng_key)注意params是 Python 层的一个 pytreepmap之后每个设备自动持有它自己的副本lax.pmean负责把所有副本的平均梯度聚合回来再让每张卡用它更新自己的参数。整个函数没有一处出现rank、device或all_reduce但实际执行时 XLA 为你插入了标准的集合通信。2.3 pjit 和自动分片一条更激进的路线pmap只能做数据并行。当你的模型大到单卡放不下要切分参数、多头注意力矩阵PyTorch 通常要手动写张量并行或引入 Megatron-LM、vLLM 这类外部库。JAX 则把切分也变成了声明式配置即pjit新版叫jax.jitout_shardings的自动分片。你可以定义一个有逻辑形状的大参数然后用PartitionSpec描述它在各设备上如何排布比如P(dp, mp)表示第一维按数据并行组切第二维按模型并行组切。XLA 会根据这个描述自动算出通信方案需要 ReduceScatter 的地方插 ReduceScatter需要 AllGather 的地方插 AllGather。这意味着你从“手写通信代码”切换到了“描述数据排布”编译器和运行时帮你剩下的活。举个具体例子假设 mesh 是二维的dp4, mp2一共 8 设备你要切一个[dim0, dim1]的权重矩阵from jax.sharding import Mesh, PartitionSpec as P from jax.experimental import mesh_utils import jax.numpy as jnp devices mesh_utils.create_device_mesh((4, 2)) mesh Mesh(devices, (dp, mp)) def init_weight(key, shape, mesh): # 在全局逻辑形状下初始化 param jax.random.normal(key, shape) # 声明分片策略 return jax.device_put(param, jax.sharding.NamedSharding(mesh, P(dp, mp)))之后对这个参数做任何矩阵乘法XLA 自动知道它需要先跨 mp 组做 AllGather 才能并行计算算完再按分片规则切回来。这套做法的上限很高因为全局优化器看到的通信模式更完整下限也低因为你一旦描述错分片编译器会给你难读的报错Debug 成本比 PyTorch 显式通信的版本高得多。2.4 两种模型并行的“心智负担”完全不同跑熟 PyTorch 张量并行的人都知道你得时刻关注“当前 GPU 上这块内存属于哪个 rank 的哪一段参数”写代码时的视角是硬件视角。JAX 的 pjit 让你回到“逻辑视角”你写的是全局模型分片只是给编译器看的注解。这个心智转换对新手来说是福利对老手来说是挑战——习惯了手写通信的人一开始很难信任“编译器能替我排好”。我在一个 GPT 规模模型上实测过参数矩阵按P(dp, mp)切分后XLA 自动生成了通信图整体吞吐比手动 Megatron 式切分大约只低 2%~5%这 5% 主要来自它做了保守的 AllGather 调度。但代码量大概只剩原来的 40%。如果模型是标准的 Transformer这个方案性价比相当高。但如果你的模型有很强的自定义结构比如稀疏计算、动态控制流自动分片可能就不是最优解编译器会保守地插入大量通信。3. 实操对照同一套模型在两边跑多卡的具体差别3.1 环境准备conda、CUDA 和第一道门槛我相信很多读者都是从 PyTorch 入门的conda 建环境、装 GPU 版、跑通 mnist这一套闭着眼睛都能搞定。JAX 的入门门槛要高一截尤其是 GPU 版安装你需要认清自己的 CUDA 版本然后去 jax 官方索引手动指定对应的jaxlib版本。例如 CUDA 12 系列pip install -U jax[cuda12]但这里有个历史坑jax[cuda12]默认会装最新版 jaxlib它要求 driver 版本至少 525 或以上老卡比如 10 系、20 系可能直接跑不起来。更稳的做法是pip install jax0.4.23 jaxlib0.4.23cuda12.cudnn89 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.htmlPyTorch 那边装完conda install pytorch pytorch-cuda12.1 -c pytorch -c nvidia就完事生态成熟度差距在这一步就体现出来了。JAX 在 pip 和 conda 的长期支持上没有那么齐全很多发行版需要自己去 GitHub Release 页面对照版本。不过一旦装好JAX 对多卡的初始化和设备发现比较简单不需要像 PyTorch 那样配置MASTER_ADDR、MASTER_PORT等环境变量当然多机场景还是要的。3.2 数据加载的差异Dataset 与纯函数式迭代器PyTorch 数据流依赖DataLoader、Sampler、DistributedSampler在分布式场景还要设置world_size和rank去做切分。JAX 没有官方 DataLoader通常用tf.data或自写numpy切片迭代器。如果纯用tf.data你需要额外 pip 安装tensorflow如果用datasets库它自带IterableDataset。我在多机数据并行时见过很多团队因为DistributedSampler用错导致一个 epoch 的样本有重复或缺失。JAX 这边没有类似的现成 Sampler你得自己在进入循环前算好per_device_batch 16 num_devices jax.device_count() global_batch per_device_batch * num_devices for step in range(total_steps): # 每次取出 global_batch 条按设备数切分 x next(iter(dataset)) x x.reshape(num_devices, per_device_batch, *x.shape[1:])表面上是少了“官方支持”的封装实际上是被迫去理解数据到底怎么分布到设备上长远看对训练稳定性是有好处的。PyTorch 虽然开箱即用但数据切分细节对你隐藏后一旦出了问题排查链路反而更长。3.3 随机数处理的画风突变PRNGKey 与 pmap 的纠缠PyTorch 里torch.manual_seed(0)是全局限定分布式训练时通常每个 rank 设不同 seed 保持多样性或者用固定 seed DistributedSampler保证可复现。JAX 没有全局种子这回事所有随机性都来自显式传入的PRNGKey而且要遵守“拆分-使用-丢弃”的原则。在pmap里做随机操作时会有个隐蔽问题每张卡需要一个不同的 key否则全部生成相同的 dropout mask。正确做法是key jax.random.PRNGKey(0) keys jax.random.split(key, num_devices) # 拆成 num_devices 份 pmap def train_step(params, x, y, key): dropout_mask jax.random.bernoulli(key, p0.5, shapex.shape) ...如果你在pmap函数内部直接调用jax.random.PRNGKey(0)XLA 会在每张卡上生成同样的随机数模型各卡之间就失去多样性。这个是所有人上手时必踩的坑。一句话总结每一个pmap内消耗随机数的操作都必须拿到一个从外部拆分好的 key。3.4 模型的“有状态”习惯要彻底改掉PyTorch 里模型是对象nn.BatchNorm里的running_mean是 bufferDropout 是训练/评估模式切换optimizer 有内部状态如 momentum。JAX 中这些统统要变成显式参数或显式输入你要保存running_mean就把它们放进一个 pytree 一起传入你要切换 dropout就把布尔开关传进函数你要带 momentum就自己写state字典维护。这乍一听很麻烦但在多卡并行下反而减少了混淆。PyTorch DDP 在某些状态下比如 train/eval 切换、buffer 同步会出现“部分卡的状态不一致”问题JAX 的显式状态天然规避了这种不一致因为所有卡上跑的同一份输入 pytree 产生相同的结果。我们团队迁移后BatchNorm 的 buffer 同步问题直线下降因为根本不存在隐式 buffer 了。4. 性能和调试的现实差异编译开销、日志与排查4.1 编译时间是一块硬伤JAX 的极致优化是从 JIT 编译换来的代价是第一个 step 特别慢。一张 H100 上跑小型 Transformer首次编译可能花费 20~60 秒如果函数里有条件分支且分支里有不同的张量形状XLA 会为每种形状各生成一个特化版本最坏情况下线性增长。PyTorch 的 Eager 模式基本没有编译时间所以你在线调试时能用 print 大法逐步查。JAX 里你一旦把函数包进jit或pmapprint 就只能在 host 端看到设备端中等张量的值没法直接看。这个差异导致两种开发习惯PyTorch 是“跑起来调试”JAX 是“想清楚再跑”。我个人的做法是先在 CPU 上、不jit的情况下用jnp单步跑通逻辑再包上jax.jit验证 shape 和 dtype最后才上pmap。每加一层变换就多做一轮 debug否则三层错误混在一起报错信息能让你怀疑人生。4.2 Debug 手段host_callback、debug.print和各种不彻底的办法JAX 官方目前推荐用jax.debug.print在设备上打印但它只在 host 端输出且打印的顺序不保证与计算顺序一致。更传统的办法是用jax.experimental.io_callback旧版host_callback把设备值传回主机处理。比如from jax.debug import print as jprint pmap def train_step(params, x, y, key): jprint(loss step {step}: {loss}, stepstep, lossloss) ...注意这会把设备端数据同步回 CPU显著拖慢训练。调试完毕后务必全部删掉。相比之下 PyTorch 的pdb/torch.distributed.barrier虽然也不优雅但至少所见即所得。如果你是从 PyTorch 转过来的我建议先不要追求“一步到位 pmap”先在单卡jit下跑通再看 pmap 的报错。4.3 Checkpoint 保存别天真地以为torch.save能平移PyTorch 的state_dicttorch.save是把参数、optimizer 状态整体序列化恢复时load_state_dict即可。JAX 没有统一的save接口常规做法是用flax.training.checkpoints.save_checkpoint如果你用 Flax或者自己用np.save/pickle把 pytree 序列化成字典再存。跨框架迁移 checkpoint 是另一件麻烦事PyTorch 的权重是torch.TensorJAX 是jnp.ndarray两者 layout 头、字节序虽然大体兼容但需要用numpy中转。比如import numpy as np import torch # 读入 PyTorch 权重 sd torch.load(model.pt, map_locationcpu) # 转 numpy 字典 np_dict {k: v.numpy() for k, v in sd.items()} # 再根据 jax/flax 的命名规则手动映射如果你只是简单把参数名改一下有机会直接复用但遇到nn.Linear的参数名weight、bias和 Flax 的kernel、bias对应关系得写映射表。这个工作虽然机械但在迁移大模型时非常耗时建议提前写好一套自动转换脚本。4.4 性能指标对比什么场景值得折腾不能光说差异得给数据参考。同样的 7B 规模 Llama 风格模型在同一个 8×H100 节点上PyTorch FSDP 全分片默认配置约 35k tokens/s 左右视实现和通信优化而定JAX pjit 自动分片dpmp 二维 mesh约 38k~40k tokens/s同时显存占用略低因为编译期能消除更多临时 buffer差距并没有传说中那么大真正拉开差距的场景是超大 batch 和自定义梯度操作。JAX 在编译器层的融合能力更强如果你反复对中间量做grad、vmap组合优化空间明显PyTorch 的优势则在于常用算子被社区的 CUDA 优化得很透eager 模式下的常规模块性能已经很接近理论峰值。另外注意PyTorch 的 torch.compileinductor也在不断缩小编译差距但它和 JAX 的 XLA 编译在分布式场景下的控制力度仍不是一个量级。5. 常见问题与排查技巧实录5.1 “pmap 输出的数组回不到 CPU”——ShardedDeviceArray 的坑pmap 函数输出的张量在 JAX 里是ShardedDeviceArray天然分布在多卡上。如果你想拿它到主机端做np.mean或画曲线必须显式jax.device_get(result)把它全量汇聚回 CPU。很多新手第一次跑通后直接print(result)会发现打印出/device:GPU:0甚至GdaGlobalDeviceArray一脸懵。我的经验是在训练循环里尽量只让 pmap 返回标量 loss并且用一个 list 反复device_get累积比如每 10 步取一次带回 CPU避免频繁全量通信干扰训练。日志里的 loss 曲线用这种方式记录就足够了。5.2jit里出现 Python 控制流——动态 shape 报错JAX 函数在编译时需要静态已知每个数组的形状。如果你写了if x.shape[0] 128: x x[:128]这在 Python 层执行XLA 需要为x.shape[0]生成一个常量特化版本。一旦 batch size 变化就会重新编译QPS 骤降。正确做法是用lax.cond或jax.lax.scan等函数式控制流原语x lax.cond(x.shape[0] 128, lambda t: t[:128], lambda t: t, x)如果 batch size 在训练中确实会变尽量预先 pad 到固定 shape否则“编译风暴”会拖垮迭代节奏。5.3 通信与计算重叠不要低估 XLA 的行为PyTorch 里all_reduce和下一轮前向可以用 stream 手动重叠很多人也习惯了这种精细控制。JAX 不暴露 streamXLA 会自己决定插卡。大多数情况下好消息是它能做得更好坏消息是你不容易干预尤其当你想做梯度累积、正交性约束等高阶操作时很难精确指定“在哪一次集合通信后插入同步”。我踩过的一个具体问题在一组自定义 comm 算子类似lax.all_gather附近做device_get触发同步后通信延迟骤增。解决办法是只允许在训练 step 最外层做host同步内部所有集体操作都不要“触摸”主机。5.4 多机注意mesh 创建与设备顺序单机多卡时jax.devices()返回设备顺序通常和 PCIe 拓扑一致但多机时需要你用jax.experimental.mesh_utils.create_device_mesh显式构建 mesh否则默认 mesh 可能按process_index排列这在跨节点做数据并行时会导致通信走主机网络而不是 NVLink吞吐下降明显。我在 2 机 8 卡上试过这个场景默认 mesh 下每 step 慢 25% 左右改成mesh mesh_utils.create_device_mesh((2, 4))让进程间维度排在最内层延迟立刻降下来。这个细节很容易被忽略因为单机跑得非常欢多机一上就“莫名其妙”慢了。5.5 无法反序列化和 float16 的精度问题JAX 默认jnp.float32但在 GPU 上训练时很多人习惯x x.astype(jnp.bfloat16)来省显存。注意pmap通信聚合时lax.pmean在 bfloat16 下的精度损失比 float32 大得多梯度归约时容易出现“最后一层权重震荡”。稳妥做法是计算梯度时保持 float32只在显存敏感的前向传播时选择性用 bfloat16。这不是框架 bug而是混合精度策略没做对。6. 选型建议什么场景该留在 PyTorch什么场景值得试试 JAX6.1 五个“留在 PyTorch”的典型信号如果你的项目满足任意两条我不建议你急着迁 JAX团队已经用 PyTorch 栈维护了大量数据管道和推理服务统一框架的成本很低重写收益抵不过迁移风险模型包含大量自定义算子、稀疏计算、动态控制流PyTorch 的 eager 模式调试最顺手你重度依赖 HuggingFace 生态transformers仓库对 PyTorch 支持最完善JAX 不是没有但部署资料少得多虽然 JAX 版本后来也有但你查 issue 时会发现很多 PyTorch 特有的 helper 用不了需要快速迭代验证研究 ideaPyTorch 的 run 起来就是快写起来心智负担小你的集群运维是成熟 PyTorch 体系比如集成 MLflow、WB、监测 agent 都按 PyTorch 生态打造的6.2 四个值得跨过学习曲线去尝试 JAX 的场景反过来说下面这几种情况JAX 的优势比较明显在做大规模语言模型训练尤其是需要自动张量并行 数据并行混合你不希望手写 Megatron 式切分代码模型里大量使用grad、vmap、scan这类变换式 API比如神经算子、RNN 定制实现JAX 的编译优化能明显加速你希望在 TPU 上跑实验不需要为 TPU 单独改太多代码JAX 是 TPU 上的第一公民你在追求极致的自定义训练管线比如深层强化学习、进化策略、多任务博弈JAX 的纯函数风格更适合做高层次的算法抽象6.3 如果决定迁移我建议走这三步不要幻想一步到位把整个代码库翻完。我亲测比较稳妥的迁移路径是第一步先把数据加载从 PyTorch DataLoader 换成tf.data或datasets的迭代器因为在 JAX 里数据管道也推荐函数式风格。第二步写一个最小可跑的 Flax或 Haiku模型把 forward 和 loss 都封装成纯函数然后在单卡上验证梯度更新、checkpoint 保存恢复。第三步再上jit和pmap/pjit这时候你才真正接触到 JAX 分布式训练的核心mesh 设计、分片策略和通信模式的声明式表达。我见过太多人第一步就上了pmap结果把环境和函数式改造的坑混在一起最后两周没跑顺得出“JAX 不成熟”的结论。其实不是框架不成熟是学习顺序不太对。先把基础变换吃透分布式层只是多一个 axis 的问题。6.4 最后说说我的真实感受这套从 PyTorch 迁到 JAX 的过程走完最大的收获其实是更深刻地理解了“分布式训练的本质是数据排布和通信规划”而不只是“多卡跑 faster”。PyTorch 让你习惯性地把通信当作分布式运行时的一个附属品而 JAX 从一开始就逼你把并行写进程序的语义里。这一点决定了天花板如果你只想安安静静做实验、跑通任务PyTorch 足够好如果你需要把模型规模再往上顶一个台阶并且愿意花时间跨过编译器的学习曲线JAX 值得一试。我自己现在的混合策略是日常小实验、模型改动频繁的研究项目仍用 PyTorch涉及 7B 以上模型、需要做大规模数据/模型并行混合的核心训练任务用 JAX。两个生态各有各的顺手之处能熟练切换比“拥护某个框架”要实用得多。