上个月把公司内部一个多卡训练任务从 PyTorch 迁移到了 JAX，起因倒不是跟风，而是模型里频繁涉及自定义梯度、大批量数据并行，以及尝试在单机多卡上跑一些张量并行的实验。折腾了大概两周，踩了不少坑，也把 JAX 分布式训练和 PyTorch 的底层差异摸得比较清楚了。这篇文章不打算写成翻译文档，我想直接聊聊两者在分布式的设计哲学上有多不一样，以及在实操中这些差异是怎么体现出来的。无论你是正在观望要不要学 JAX，还是已经在 PyTorch 里用 DDP 跑惯了多卡，这篇文章都值得你花十分钟过一遍——它至少能帮你想明白一个问题：换框架，到底换的是 API 还是思维方式？
1. JAX 和 PyTorch 在底层哲学上的分岔路
1.1 JAX 本质上是“可微分的 NumPy”，不是“另一个 PyTorch”
很多人第一次打开 JAX 文档，会下意识拿它和 PyTorch 对标，然后发现它连nn.Module都没有，第一反应是“这玩意儿能写模型吗”。其实 JAX 对自己的定位从来不是一个轻量级深度学习框架，而是一套基于 XLA 的数组计算与自动微分系统，它的核心组件是jnp（NumPy API）、grad（自动微分）、jit（即时编译）和vmap（向量化映射）。模型只是在这些基础组件之上叠加出来的“用法”，而不是框架内置的抽象。
这个定位差异直接影响了分布式训练的形态。PyTorch 的思维是“命令式 + 模块化”，模型是一个有状态的对象，分布式训练里的通信、梯度同步、参数广播都是围绕这些对象做的。JAX 的思维是“函数式 + 变换”，模型只是一个纯函数params -> output，分布式并行本身也是函数变换的一种。所以你看 JAX 文档里几乎不会出现“训练循环”这种词，它一直在强调的是jit、vmap、pmap、pjit这些变换怎么组合。
这个区别说起来抽象，落到代码上就非常直观。PyTorch 里你要想给某个模块做参数平均，得显式写出model.parameters()遍历、all_reduce之后除以世界大小；JAX 里你只需要在pmap里写p.mean(params, axis_name="batch")，通信自动产生。因为 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_name，XLA 就会自动插入对应的通信原语。所以从用户视角看，分布式训练和多设备之间的通信几乎消失在了函数语义里。
来看一个直观对比。同样是 8 卡数据并行训练，PyTorch 的启动方式通常是：
# PyTorch DDP import torch.distributed as dist dist.init_process_group(backend="nccl") 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_device=local_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_name="batch") # 自动跨卡归约 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 层的一个 pytree，pmap之后每个设备自动持有它自己的副本，lax.pmean负责把所有副本的平均梯度聚合回来，再让每张卡用它更新自己的参数。整个函数没有一处出现rank、device或all_reduce，但实际执行时 XLA 为你插入了标准的集合通信。
2.3 pjit 和自动分片：一条更激进的路线
pmap只能做数据并行。当你的模型大到单卡放不下，要切分参数、多头注意力矩阵，PyTorch 通常要手动写张量并行或引入 Megatron-LM、vLLM 这类外部库。JAX 则把切分也变成了声明式配置，即pjit（新版叫jax.jit+out_shardings的自动分片）。
你可以定义一个有逻辑形状的大参数，然后用PartitionSpec描述它在各设备上如何排布，比如P("dp", "mp")表示第一维按数据并行组切，第二维按模型并行组切。XLA 会根据这个描述自动算出通信方案：需要 ReduceScatter 的地方插 ReduceScatter，需要 AllGather 的地方插 AllGather。这意味着你从“手写通信代码”切换到了“描述数据排布”，编译器和运行时帮你剩下的活。
举个具体例子，假设 mesh 是二维的，dp=4, mp=2，一共 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 jax==0.4.23 jaxlib==0.4.23+cuda12.cudnn89 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.htmlPyTorch 那边装完conda install pytorch pytorch-cuda=12.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, p=0.5, shape=x.shape) ...如果你在pmap函数内部直接调用jax.random.PRNGKey(0)，XLA 会在每张卡上生成同样的随机数，模型各卡之间就失去多样性。这个是所有人上手时必踩的坑。一句话总结：每一个pmap内消耗随机数的操作，都必须拿到一个从外部拆分好的 key。
3.4 模型的“有状态”习惯要彻底改掉
PyTorch 里模型是对象，nn.BatchNorm里的running_mean是 buffer，Dropout 是训练/评估模式切换，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或pmap，print 就只能在 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}", step=step, loss=loss) ...注意这会把设备端数据同步回 CPU，显著拖慢训练。调试完毕后务必全部删掉。相比之下 PyTorch 的pdb/torch.distributed.barrier虽然也不优雅，但至少所见即所得。如果你是从 PyTorch 转过来的，我建议先不要追求“一步到位 pmap”，先在单卡jit下跑通，再看 pmap 的报错。
4.3 Checkpoint 保存：别天真地以为torch.save能平移
PyTorch 的state_dict+torch.save是把参数、optimizer 状态整体序列化，恢复时load_state_dict即可。JAX 没有统一的save接口，常规做法是用flax.training.checkpoints.save_checkpoint（如果你用 Flax），或者自己用np.save/pickle把 pytree 序列化成字典再存。
跨框架迁移 checkpoint 是另一件麻烦事：PyTorch 的权重是torch.Tensor，JAX 是jnp.ndarray，两者 layout 头、字节序虽然大体兼容，但需要用numpy中转。比如：
import numpy as np import torch # 读入 PyTorch 权重 sd = torch.load("model.pt", map_location="cpu") # 转 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 自动分片（dp+mp 二维 mesh）：约 38k~40k tokens/s，同时显存占用略低，因为编译期能消除更多临时 buffer
差距并没有传说中那么大，真正拉开差距的场景是超大 batch 和自定义梯度操作。JAX 在编译器层的融合能力更强，如果你反复对中间量做grad、vmap组合，优化空间明显；PyTorch 的优势则在于常用算子被社区的 CUDA 优化得很透，eager 模式下的常规模块性能已经很接近理论峰值。另外注意，PyTorch 的 torch.compile（inductor）也在不断缩小编译差距，但它和 JAX 的 XLA 编译在分布式场景下的控制力度仍不是一个量级。
5. 常见问题与排查技巧实录
5.1 “pmap 输出的数组回不到 CPU”——ShardedDeviceArray 的坑
pmap 函数输出的张量在 JAX 里是ShardedDeviceArray，天然分布在多卡上。如果你想拿它到主机端做np.mean或画曲线，必须显式jax.device_get(result)把它全量汇聚回 CPU。很多新手第一次跑通后直接print(result)，会发现打印出/device:GPU:0甚至Gda（GlobalDeviceArray），一脸懵。
我的经验是：在训练循环里尽量只让 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 不暴露 stream，XLA 会自己决定插卡。大多数情况下好消息是它能做得更好，坏消息是你不容易干预，尤其当你想做梯度累积、正交性约束等高阶操作时，很难精确指定“在哪一次集合通信后插入同步”。
我踩过的一个具体问题：在一组自定义 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，什么场景值得试试 JAX
6.1 五个“留在 PyTorch”的典型信号
如果你的项目满足任意两条，我不建议你急着迁 JAX：
- 团队已经用 PyTorch 栈维护了大量数据管道和推理服务，统一框架的成本很低，重写收益抵不过迁移风险
- 模型包含大量自定义算子、稀疏计算、动态控制流，PyTorch 的 eager 模式调试最顺手
- 你重度依赖 HuggingFace 生态，
transformers仓库对 PyTorch 支持最完善，JAX 不是没有，但部署资料少得多（虽然 JAX 版本后来也有，但你查 issue 时会发现很多 PyTorch 特有的 helper 用不了） - 需要快速迭代验证研究 idea，PyTorch 的 run 起来就是快，写起来心智负担小
- 你的集群运维是成熟 PyTorch 体系，比如集成 MLflow、W&B、监测 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。两个生态各有各的顺手之处，能熟练切换比“拥护某个框架”要实用得多。
