JAX与PyTorch分布式训练对比：从DDP到XLA分片实战

1. 从一次多卡踩坑说起：为什么我要对比 JAX 和 PyTorch

去年年底接手一个 8 卡 A100 的节点，任务是把一个 7B 级别的 Transformer 从单卡跑到多卡。团队里大部分人平时写 PyTorch，我自己之前也主要用 PyTorch 做实验，但这次因为项目上游的代码库是 JAX 写的，我被迫把两套分布式训练的路子都摸了一遍。结果发现，很多人对 JAX 分布式训练的理解停留在“它也能多卡”这个层面，但真正落到代码上，JAX 和 PyTorch 的思维方式差异非常大——不是 API 换了个名字那么简单，而是从“你怎么描述并行”这个根上就不一样。

这篇东西不是教程，也不是官方文档的翻译。我想做的是把我在实际项目里踩过的坑、对比过的写法、以及最后为什么在某些场景下选了 JAX、某些场景下又退回 PyTorch 这件事讲清楚。如果你正在纠结要不要从 PyTorch 迁到 JAX，或者你已经在用 JAX 但被pmap、shard_map、pjit这些概念绕晕了，那这篇应该能帮你省不少时间。核心关键词就几个：JAX、PyTorch、分布式训练、XLA、DDP，我会围绕这几个词把两套体系的差异拆开讲。

先说结论性的判断，方便你决定要不要往下读：PyTorch 的分布式训练是“在 Python 层面把多个进程/多张卡组织起来”，DDP 是它的典型代表；而 JAX 的分布式训练是“在编译层面把计算图切分到多设备上”，XLA 是它的底座。这个根本差异决定了你写代码的方式、调试的方式、以及性能调优的切入点完全不同。下面我分几个层面展开。

2. 核心思路拆解：进程并行 vs 编译期切分

2.1 PyTorch DDP 的本质是“多进程 + 梯度同步”

PyTorch 的 DDP（DistributedDataParallel）思路很直白：你启动 N 个进程，每个进程绑定一张卡，每个进程持有一份完整的模型副本，喂不同的数据 batch，反向传播时通过 all-reduce 把梯度同步一下。模型本身没有被切分，切分的是数据。

这种模式的好处是心智负担低。你写单卡代码，然后套一层DistributedSampler，再用DistributedDataParallel包一下模型，基本就能跑。通信发生在反向传播的梯度上，NCCL 负责 all-reduce。代码大概长这样：

import torch import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP 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]) for batch in dataloader: outputs = model(batch) loss = criterion(outputs, targets) loss.backward() optimizer.step()

这里的关键点是：模型在每个进程里都是完整的。如果你有 8 张卡，就有 8 份完整的模型副本。对于 7B 模型，每张卡要放下全部参数、梯度、优化器状态，显存压力很大。所以 PyTorch 后来有了 FSDP（Fully Sharded Data Parallel），把参数也切了，但那是另一个话题。

2.2 JAX 的分布式是“单程序多数据 + 编译期分片”

JAX 的哲学完全不同。你写的是一段“单程序”，然后用jax.pmap或者jax.jit配合Mesh和PartitionSpec告诉编译器：这段计算要怎么分布到多设备上。JAX 不在 Python 层面启动多个进程来管理并行，而是把并行信息编码进计算图，交给 XLA 编译。

最直观的对比：PyTorch DDP 里你手动管理进程和数据分发；JAX 里你描述“数组怎么切”，剩下的交给编译器。比如pmap的经典写法：

import jax import jax.numpy as jnp @jax.pmap def forward(params, batch): return model_apply(params, batch) # 输入的第一维自动映射到设备 params = jax.tree_map(lambda x: jnp.stack([x] * n_devices), params) outputs = forward(params, batch)

pmap会把函数编译成在多个设备上并行执行的版本，输入数组的第一维对应设备维度。这看起来简单，但一旦涉及参数分片、通信模式、梯度聚合，复杂度就上来了。

2.3 为什么这个差异很重要

这个差异直接决定了三件事：

第一，显存模型不同。PyTorch DDP 每卡一份完整模型，JAX 可以通过PartitionSpec把参数切到不同设备上，天然支持模型并行。这意味着在同样的硬件上，JAX 更容易跑大模型。

第二，调试方式不同。PyTorch 你可以print每个进程的中间结果，用torch.distributed的调试工具；JAX 因为编译期优化，很多中间状态你看不到，调试要靠jax.debug.print和jax.make_jaxpr去看计算图。

第三，性能调优的切入点不同。PyTorch 调 DDP 主要看通信重叠、bucket size、gradient accumulation；JAX 调的是分片策略、通信原语（all-gather、reduce-scatter）、以及 XLA 的编译选项。

3. 核心细节解析：从 pmap 到 shard_map 的演进

3.1 pmap 的局限与 pjit 的登场

pmap是 JAX 最早的并行原语，它要求你显式地把数据映射到设备上，而且只支持数据并行这一种模式。它的编译单位是“每个设备一份”，所以如果你想把模型参数也切分，pmap就力不从心了。

后来 JAX 引入了pjit（现在统一到jax.jit配合in_shardings和out_shardings），它允许你用Mesh和PartitionSpec描述任意维度的分片。比如：

from jax.sharding import Mesh, PartitionSpec as P from jax.experimental import mesh_utils devices = mesh_utils.create_device_mesh((8,)) mesh = Mesh(devices, axis_names=('data',)) @jax.jit def train_step(params, batch): ... # 用 with mesh 上下文指定分片 with mesh: params = jax.tree_map( lambda x: jax.device_put(x, P('data')), params ) ...

这里P('data')表示这个数组沿data轴切分。XLA 会根据这个描述自动插入通信操作。这比pmap灵活得多，但也更容易写错——分片策略写错了，要么编译报错，要么性能暴跌。

3.2 shard_map：更底层的控制

shard_map是 JAX 后来推出的一个更底层的 API，它让你在“每个分片”的视角写代码，然后显式地调用通信原语。这有点像 MPI 的风格：

from jax.experimental.shard_map import shard_map from jax.sharding import PartitionSpec as P @shard_map( mesh=mesh, in_specs=(P('data'), P('data')), out_specs=P('data'), check_rep=False ) def compute(params_shard, batch_shard): # 这里写的是单个分片上的逻辑 # 需要通信时显式调用 jax.lax.all_gather 等 ...

shard_map的好处是控制力强，你能精确知道通信发生在哪里；坏处是心智负担大，你得自己管理分片间的数据流动。我在实际项目里，只有在pjit的自动分片效果不理想时才会降到shard_map这一层。

3.3 PyTorch 这边的对应演进

PyTorch 也不是只有 DDP。FSDP 把参数、梯度、优化器状态都切分，思路和 JAX 的分片有相似之处，但实现层面还是“多进程 + 通信”那一套。PyTorch 2.0 之后有了torch.compile，开始引入编译期优化，但分布式这块的核心还是进程组和集合通信。

一个有意思的对比是：PyTorch 的 FSDP 需要你显式地 wrap 模型层，指定sharding_strategy；JAX 的分片是声明式的，你描述“怎么切”，编译器决定“怎么通信”。前者更可控，后者更省心但更难调。

4. 实操过程：同一任务在两套框架下的实现对比

4.1 任务设定

假设我们要训练一个两层 MLP，输入维度 784，隐藏层 512，输出 10，用 8 张卡做数据并行。数据量假设是 64000 条，batch size 全局 256，每卡 32。

4.2 PyTorch DDP 实现

import os import torch import torch.nn as nn import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import DataLoader, DistributedSampler class MLP(nn.Module): def __init__(self): super().__init__() self.fc1 = nn.Linear(784, 512) self.fc2 = nn.Linear(512, 10) def forward(self, x): x = torch.relu(self.fc1(x)) return self.fc2(x) def main(): dist.init_process_group(backend="nccl") local_rank = int(os.environ["LOCAL_RANK"]) torch.cuda.set_device(local_rank) model = MLP().to(local_rank) model = DDP(model, device_ids=[local_rank]) dataset = MyDataset() sampler = DistributedSampler(dataset) loader = DataLoader(dataset, batch_size=32, sampler=sampler) optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) for epoch in range(10): sampler.set_epoch(epoch) for x, y in loader: x, y = x.to(local_rank), y.to(local_rank) optimizer.zero_grad() out = model(x) loss = nn.functional.cross_entropy(out, y) loss.backward() optimizer.step() if __name__ == "__main__": main()

启动命令：

torchrun --nproc_per_node=8 train_ddp.py

这里的关键操作点：DistributedSampler保证每个进程拿到不同的数据子集；DDP包装后，反向传播时自动 all-reduce 梯度；sampler.set_epoch(epoch)保证每个 epoch 的 shuffle 不同。

4.3 JAX 实现（用 pjit + Mesh）

import jax import jax.numpy as jnp from jax import random from jax.sharding import Mesh, PartitionSpec as P from jax.experimental import mesh_utils def init_params(key): k1, k2 = random.split(key) w1 = random.normal(k1, (784, 512)) * 0.01 b1 = jnp.zeros(512) w2 = random.normal(k2, (512, 10)) * 0.01 b2 = jnp.zeros(10) return {'w1': w1, 'b1': b1, 'w2': w2, 'b2': b2} def forward(params, x): h = jnp.maximum(0, x @ params['w1'] + params['b1']) return h @ params['w2'] + params['b2'] def loss_fn(params, x, y): logits = forward(params, x) return jnp.mean(-jnp.sum(y * jax.nn.log_softmax(logits), axis=-1)) @jax.jit def train_step(params, x, y, lr=1e-3): loss, grads = jax.value_and_grad(loss_fn)(params, x, y) params = jax.tree_map(lambda p, g: p - lr * g, params, grads) return params, loss def main(): devices = mesh_utils.create_device_mesh((8,)) mesh = Mesh(devices, axis_names=('data',)) key = random.PRNGKey(0) params = init_params(key) # 复制参数到所有设备 params = jax.tree_map( lambda x: jax.device_put(x, P()), params ) # 数据分片 x = jnp.ones((256, 784)) y = jnp.eye(10)[jnp.zeros(256, dtype=jnp.int32)] with mesh: x = jax.device_put(x, P('data')) y = jax.device_put(y, P('data')) for step in range(100): params, loss = train_step(params, x, y) if step % 10 == 0: print(f"step {step}, loss {loss}") if __name__ == "__main__": main()

注意这里的差异：JAX 里没有“进程”的概念，你写的是一个程序，Mesh和P('data')告诉编译器怎么切。jax.jit会把train_step编译成在 8 个设备上执行的版本，梯度同步由 XLA 自动插入。

4.4 参数选择与计算过程

batch size 的选择：全局 256，8 卡，每卡 32。这个数字不是随便定的。假设每卡显存 40GB，模型参数加梯度加优化器状态大概占 2GB，激活值占大头。32 的 per-device batch 在 512 隐藏层下，激活值大概几百 MB，留足余量。如果你把 per-device batch 调到 64，激活值翻倍，可能就 OOM 了。

学习率：全局 batch 256 比单卡 32 大了 8 倍，学习率通常要相应放大。经验公式是lr_new = lr_base * sqrt(global_batch / base_batch)，但这不是铁律，实际要看优化器和任务。我用的是 1e-3，比单卡的 3e-4 大一些，但没有线性放大 8 倍，因为线性放大在训练后期容易震荡。

通信量估算：MLP 参数量约 784512 + 51210 ≈ 40 万，梯度 all-reduce 每次通信 40 万 * 4 字节 ≈ 1.6MB。8 卡 all-reduce 的通信量约 1.6MB * 2 * (8-1)/8 ≈ 2.8MB。这个量很小，通信不是瓶颈。如果是 7B 模型，参数量 70 亿，梯度 28GB，通信就是大问题，这时候 JAX 的分片优势才明显。

5. 常见问题与排查技巧实录

5.1 PyTorch DDP 的典型坑

坑一：忘记sampler.set_epoch(epoch)。这会导致每个 epoch 的数据顺序一样，模型收敛变差。我见过有人训练了三天才发现这个问题，loss 曲线看着正常但验证集不涨。

坑二：DistributedSampler的drop_last问题。如果数据集大小不能被卡数整除，最后一个 batch 可能在不同进程里大小不一致，导致 all-reduce 卡住。解决办法是设drop_last=True或者手动补齐。

坑三：NCCL 超时。默认 NCCL 超时是 30 分钟，如果某个进程卡住，其他进程会等超时然后报错。调试时可以把超时调短，比如export NCCL_TIMEOUT=60，快速暴露问题。

坑四：显存不均衡。DDP 每卡一份完整模型，如果某张卡上还有其他进程占显存，就会 OOM。用nvidia-smi确认每张卡干净。

5.2 JAX 分布式的典型坑

坑一：分片策略写错导致隐式通信。比如你把参数标成P()（复制），但输入标成P('data')，XLA 可能会插入 all-gather 把参数复制到所有设备，性能暴跌。用jax.make_jaxpr看编译后的计算图能发现这类问题。

坑二：pmap和pjit混用。这两个的并行模型不一样，混用容易出奇怪的问题。新项目建议直接用pjit+Mesh，pmap只在维护老代码时用。

坑三：随机数处理。JAX 的随机数是显式的，多设备下每个设备需要不同的 key。如果你在所有设备上用同一个 key 初始化参数，那所有副本都一样，数据并行就退化成单卡了。正确做法是用jax.random.split按设备数切分 key。

坑四：编译时间过长。JAX 的编译是“一次编译，多次执行”，但如果你的输入 shape 经常变，就会反复编译。固定 shape 能大幅减少编译时间。

5.3 问题速查表

6. 工具选型与场景建议

6.1 什么时候选 PyTorch

如果你的团队已经有一套成熟的 PyTorch 训练流程，模型规模在单卡能放下的范围内，那 DDP 是最省事的选择。生态成熟，调试工具多，遇到问题网上能搜到答案。特别是做研究、快速迭代实验，PyTorch 的灵活性和可调试性是优势。

另外，如果你的模型有大量动态控制流、或者需要频繁改变计算图结构，PyTorch 的 eager 模式更友好。JAX 的 jit 对动态 shape 和动态控制流支持有限，虽然jax.lax.while_loop能处理一部分，但写起来别扭。

6.2 什么时候选 JAX

模型大到单卡放不下，需要模型并行加数据并行混合，JAX 的分片抽象更自然。Google 内部的大模型训练基本都用 JAX，不是没有道理的。另外，如果你需要极致性能，XLA 的算子融合和编译优化在特定场景下能比 PyTorch 快不少，尤其是 TPU 上。

还有一种场景是“函数式编程风格”的团队。JAX 的纯函数、无副作用、显式随机数这些设计，对于喜欢函数式风格的工程师来说很顺手。但这也意味着你得改变写代码的习惯，不能随便在函数里改全局状态。

6.3 混合使用的可能性

实际项目里，我见过用 JAX 做训练、用 PyTorch 做推理的，也见过反过来。两者之间的模型转换可以通过 ONNX 或者手动导出权重。但说实话，混合使用的维护成本不低，除非有明确的性能或生态理由，否则建议统一到一套框架。

7. 我个人的实操体会

踩了这么多坑，我最深的体会是：JAX 的分布式训练不是“PyTorch 的替代品”，而是另一种思维方式。你用 PyTorch 的思维去写 JAX，会觉得处处别扭；但一旦接受了“描述分片、交给编译器”这个范式，很多在 PyTorch 里要手动处理的事情就自动了。

另一个体会是，调试 JAX 分布式问题，计算图是你的朋友。jax.make_jaxpr和jax.jit(...).lower(...).compile()能让你看到编译后的 HLO，通信操作一目了然。PyTorch 这边，torch.distributed的 debug 工具和 NCCL 的日志是主要手段。

最后分享一个小技巧：如果你在 JAX 里不确定分片策略对不对，先用小规模数据在单机多卡上跑一遍，用jax.debug.print打印每个分片上的 shape，确认无误再上大规模。这个习惯帮我省了很多次重新编译和重新排队的时间。