“最近被朋友问到最多的问题就是‘JAX 分布式训练到底和 PyTorch 哪里不一样’。花了两周时间把一套在 PyTorch 上跑顺的多机训练流程迁到 JAX 上之后我真的觉得这两套框架的‘分布式’完全是两种生物。PyTorch 的 DDP 是显式的命令式玩法——每个进程写同一份脚本手动初始化进程组自己处理 DataLoader 的分布式采样跑完 backward 再同步梯度而 JAX 的 pmap / pjit 把‘多设备并行’做成了编译器的一部分你用一套函数式变换就能表达整台机器、整个集群的并行逻辑。这篇文章不打算从 hello world 开始讲而是直接对比两者在分布式训练上的本质差异、踩坑细节以及一份可以直接照着跑的 JAX 多卡数据并行代码。”1. JAX 分布式训练和 PyTorch DDP 的本质差异先看编程范式1.1 XLA 编译器如何把“多设备”变成计算图的一部分很多从 PyTorch 过来的朋友第一次看 JAX 的代码会有点懵这不就是 numpy 吗确实是 numpy 的语法但背后藏着一个 XLA 编译器。当你调用jax.jit或jax.pmap时你写的 Python 函数不会立刻逐行执行而是会被 JAX 先“trace”成一张计算图再交给 XLA 去编译。编译过程中XLA 会做算子融合、内存复用、数据布局优化甚至会把跨设备通信也当成计算图里的普通算子来处理。PyTorch 默认是 eager 模式一行代码对应一次真实的算子执行整个计算图是“边跑边建”的。分布式训练对它来说更像是一个外挂你用torch.distributed.init_process_group启动一个进程组再用DistributedDataParallel把模型包一层之后 DDP 在 forward 时注册 autograd hookbackward 时自动做梯度 allreduce。这套机制本身设计得很精巧但它始终是“挂在框架外面”的额外层而不是编译器原生能力。JAX 的思路完全不同你告诉它“我要把这个函数跑在 8 张卡上”XLA 在生成设备代码时就直接考虑数据分片、通信顺序、通信与计算的重叠。所以 JAX 里的分布式不是某个 API而是函数变换pmap带来的编译结果。这个差别带来的直接体验是JAX 的多卡代码往往比 PyTorch 干净很多你不需要关心rank、local_rank这些进程身份因为绝大多数时候你写的还是一份“普通”的逻辑。1.2 不可变参数与纯函数训练循环JAX 强烈推崇函数式编程模型参数在训练过程中不是“原地更新”的而是每次更新后返回一份新的参数结构。你在 PyTorch 里习惯的param.grad、optimizer.step()这套可变状态玩法在 JAX 里几乎不存在。通常的做法是把参数放在一个 dict 或者更一般的 pytree 结构里带optax优化器算完更新后训练函数返回“新的参数 dict”。这个看似反人类的设计在分布式训练里反而是个大优势。纯函数没有副作用同一个函数可以被安全地复制到 N 张卡上分别执行设备之间只需要同步梯度不需要像 PyTorch DDP 那样偷偷修改模型的内部状态。pmap的语义也因此特别清楚每个设备拿到输入的局部切片跑同样的纯函数输出再拼回来。你用不着担心某个设备上的随机状态污染了另一个设备也不用担心反向传播的 hook 顺序出错因为这些在函数式世界里天然不存在。第一次写 JAX 训练循环时我甚至因为不习惯“参数更新必须返回新对象”而犯了低级错误直接在 dict 里做减法导致pmap编译时报错。后来养成习惯所有参数的读写都走函数参数和返回值训练逻辑一下子就清爽了。1.3 单程序多数据 vs 多程序多数据PyTorch DDP 属于典型的多程序多数据MPMD模式。你用torchrun --nproc_per_node8 train.py启动 8 个独立进程每个进程独立加载模型、独立跑数据靠显式的allreduce通信保证梯度一致。进程之间通过rank区分身份代码里到处是if rank 0之类的判断。JAX 的pmap属于单程序多数据SPMD模式。你只写一份逻辑JAX 在底层把这份逻辑复制到多个设备上执行程序员不需要手动管理进程身份和通信原语。真正需要多机协同的时候你也只需要调用一次jax.distributed.initialize()之后pmap自动感知全局设备。这个心智模型差异特别重要。PyTorch 的报错信息里经常出现rank、socket、world_size初学时容易一头雾水JAX 则把这些概念藏起来了你面对的是更纯粹的“设备阵列”。当然藏起来不等于没有理解 SPMD 背后的通信机制依然必要尤其是你要做模型并行的时候。2. JAX 分布式训练三件套initialize、pmap、pjit 的实操拆解2.1 用 jax.distributed.initialize 初始化集群在很长一段时间里JAX 的多机初始化都依赖环境变量和命令行参数用起来比较别扭。从 0.4.1 版本开始官方推荐在代码里直接调用初始化函数。单机多卡时你甚至可以不调用任何分布式初始化函数pmap会直接使用所有本地设备。多机多卡时每台机器都要调用一次jax.distributed.initialize()第一个进程承担 coordinator 角色其他进程通过它建立连接。import jax # 协调节点机器 0 jax.distributed.initialize( coordinator_address10.0.0.1:1234, num_processes2, process_id0, ) # 机器 1 jax.distributed.initialize( coordinator_address10.0.0.1:1234, num_processes2, process_id1, )眼尖的朋友会发现这段代码里的process_id需要你手动区分。实际项目中我习惯用环境变量传入再写个小函数统一初始化import os import jax coordinator os.environ.get(COORDINATOR_ADDRESS, localhost:1234) rank int(os.environ.get(RANK, 0)) world_size int(os.environ.get(WORLD_SIZE, 1)) jax.distributed.initialize( coordinator_addresscoordinator, num_processesworld_size, process_idrank, )这样就能和 Kubernetes 的 Job、SLURM 这类任务调度工具无缝对接。RANK、WORLD_SIZE都是调度器常注入的环境变量迁移成本很低。2.2 pmap 实现数据并行从单卡到多卡只改三处pmap是数据并行的主力 API。名字里的p是 parallel它的作用是把一个函数映射到多个设备上并且自动沿函数的第一个参数维度切分数据。最简单的一句话总结你把原来要跑的 batch 从(batch_size, ...)改成(num_devices, per_device_batch, ...)然后用pmap包装训练函数JAX 就会像map一样把每个局部 batch 送到对应设备各自计算梯度再自动做跨设备的梯度同步。from jax import pmap, lax, grad, device_put_replicated, local_devices import optax def loss_fn(params, x, y): pred x params[w] params[b] return jnp.mean((pred - y) ** 2) def train_step(params, batch): x, y batch grads grad(loss_fn)(params, x, y) # 关键必须手动做跨设备平均pmap 不会自动做 allreduce grads lax.pmean(grads, axis_namedevices) return optax.apply_updates(params, grads) # 单卡逻辑先验证 params {w: jnp.zeros((16, 1)), b: jnp.zeros((1,))} train_step_single pmap(train_step, axis_namedevices) # 如果只有一张卡也能跑 # 多卡先把参数复制到所有设备 devices local_devices() params device_put_replicated(params, devices) # 训练时把 batch 切分到设备维度 x x.reshape(len(devices), -1, x.shape[-1]) y y.reshape(len(devices), -1, y.shape[-1]) params train_step(params, (x, y))从单卡到多卡你只需要改三处用device_put_replicated复制参数、把 batch reshape 成设备数维度、在训练函数里给梯度加lax.pmean。比起 DDP 要配置 init_process_group、DistributedSampler、DistributedDataParallel 三个组件确实省了一截。2.3 需要模型并行时 pjit 怎么切pmap擅长数据并行但当模型大到单卡放不下时就需要模型并行。JAX 的答案是pjit在新版本里也逐步并入jax.jit的 sharding 体系它允许你通过sharding规则告诉编译器“哪些参数放在哪些设备上”剩下的切分、通信、合并全部由 XLA 自动完成。from jax.experimental.pjit import pjit from jax.sharding import Mesh, PartitionSpec mesh Mesh(devices, (data, model)) sharding PartitionSpec(data, None, None, None) # 第一个维度按数据并行切 model_step pjit( train_step, in_shardings({}, sharding, sharding), out_shardings({}, None), )这段代码的意思是参数 dict 的第一个字段比如w按数据维度切b不切batch 按数据维度切。实际跑大模型时你会把部分参数按model维度切到不同设备上PyTorch 那边就得用 FSDP 手动配置sharding_strategy、auto_wrap_policy这一堆东西。这里不否定 FSDP 的价值只是想说 JAX 把 sharding 的描述抬到了声明式层面你告诉编译器“我想怎么放”剩下的是编译器的工作。3. 与 PyTorch DDP 对照通信、调试、生态的差距3.1 梯度同步的真实发生位置PyTorch DDP 会在模型初始化时把参数注册为 autograd 图的叶子节点backward 过程中每个参数的梯度算完立即触发一次针对这个梯度桶的 allreduce。由于通信是异步的多算子可以重叠整体效率很高。JAX 的梯度同步发生在你的函数式代码里——lax.pmean或lax.psum。当pmap带着axis_name编译这段代码时XLA 会把pmean识别成跨设备 allreduce 操作并自动安排通信与计算的依赖顺序。两者在数学上等价但有一处细节非常容易踩坑PyTorch DDP 默认帮你把梯度平均了而 JAX 的pmap并不会自动平均梯度——你必须在训练函数里显式写lax.pmean(grads, axis_namedevices)。我第一次跑多卡 JAX 训练时漏了这行结果每张卡都在各自优化模型一直不收敛排查了大半天。这是从 PyTorch 迁过来的人最常见的坑后文我还会再强调。3.2 调试习惯的天壤之别PyTorch eager 模式的调试体验是无可挑剔的你可以随时print(tensor.shape)、打断点、在forward里写一堆日志一切都是“看得到摸得着”的。JAX 一旦进入jit或pmap你的函数就被编译成了一整张 XLA 计算图在函数内部直接print不会按你的直觉逐行执行打印出来的多半是Traced或者编译信息而不是具体数值。我的调试方法是分三步走先用纯 eager 模式不写jit、pmap把训练逻辑跑通验证 loss 在下降再用jax.debug.print代替print在编译函数里输出中间量最后才加jit或pmap并且每次只改一个变量。jax.debug.print是个很好用的官方调试工具它是为“在编译图里打印真实值”设计的。虽然比 PyTorch 的print绕但习惯之后效率也不算差。3.3 生态对比PyTorch 生态 vs Flax/Optax 生态PyTorch 的生态地位目前依然不可撼动研究社区几乎所有前沿模型、层实现、预训练权重都原生支持 PyTorch遇到问题随便一搜就是答案。JAX 的生态主要集中在 Flax、Optax、Elegy 以及 HuggingFace 的 Flax 分支上这几年明显在增长但总体体量仍然不如 PyTorch。不过分布式训练这个场景JAX 的优势反而被很多大厂团队认可。OpenAI 早期不少大规模训练用的就是 JAXDeepMind 更是 JAX 的重度用户。如果你做的是大模型预训练、RL 这类对横向扩展和编译器优化敏感的工作JAX 的干净抽象和自动分片能力确实能节省不少团队协作成本。我做多机训练时最直观的感受是 JAX 代码的行数比 PyTorch 少很多不需要写 DDP 的样板代码不需要处理DistributedSampler甚至不需要手动计算梯度同步的时机。我把常用的对比项放在一张表里方便大家快速建立坐标系对比项PyTorchDDP / FSDPJAXpmap / pjit编程范式命令式动态图函数式编译图并行模型多程序多数据MPMD单程序多数据SPMD参数存储模型内部可变状态函数外的不可变 pytree梯度同步DDP 自动注册 hook需要显式lax.pmean模型并行FSDP / 手动 shardingpjit / shard_map 自动切分调试体验断点、print 直接可用eager 模式或jax.debug.print生态成熟度最成熟覆盖面广Flax/Optax 生态成长中多机多卡初始化init_process_grouptorchrunjax.distributed.initialize()4. 实战演练用 JAX 跑一个分布式线性回归4.1 环境准备JAX 的 GPU 安装要求JAX 的环境相对 PyTorch 简单一些但要养成先查版本再装的习惯。CPU 版本直接安装pip install -U jax jaxlibGPU 版本推荐用 JAX 官方支持的 CUDA 12 预编译包pip install -U jax[cuda12]这里特别提醒JAX 对 GPU 驱动和 CUDA 工具链的依赖比 PyTorch 更“挑剔”一些因为 XLA 需要把自己的运行时和驱动兼容起来。如果你在 Windows 上跑 GPU官方目前更推荐用 WSL2 或者 DockerWindows 原生支持不是没有但坑比较多。我自己的项目就是在 WSL2 里跑的稳定性很好。安装完之后用下面几行检查设备和分布式环境是否正常import jax print(jax.devices()) print(jax.local_device_count()) jax.distributed.initialize() print(jax.device_count())如果jax.devices()只看到 CPU 而看不到 GPU多半是 jaxlib 的 CUDA 版本和本机驱动不匹配或者 GPU 驱动没有正确暴露给 XLA。可以先跑nvidia-smi确认驱动状态再确认安装的是jax[cuda12]而非纯 CPU 版本。4.2 完整训练脚本逐段拆解我们用一个最简的线性回归做演示目标不是精度而是把pmap的分布式气质展示清楚。import jax import jax.numpy as jnp import optax from jax import grad, pmap, lax, device_put_replicated, local_devices # 1. 定义损失函数 def loss_fn(params, x, y): pred x params[w] params[b] return jnp.mean((pred - y) ** 2) # 2. 定义单步更新这里先不写 pmap方便单卡验证 def train_step(params, batch): x, y batch grads grad(loss_fn)(params, x, y) grads lax.pmean(grads, axis_namedevices) return optax.apply_updates(params, grads) # 3. 用 pmap 包成多卡版本 train_step pmap(train_step, axis_namedevices) # 4. 初始化参数先放在 CPU 上再复制到所有设备 params {w: jnp.zeros((16, 1)), b: jnp.zeros((1,))} devices local_devices() params device_put_replicated(params, devices) # 5. 生成训练数据 def make_batch(key, batch_size, dim16): x jax.random.normal(key, (batch_size, dim)) w_true jax.random.normal(key, (dim, 1)) y x w_true 0.1 * jax.random.normal(key, (batch_size, 1)) return x, y # 6. 训练循环 key jax.random.PRNGKey(0) num_devices len(devices) per_device_bs 64 global_bs per_device_bs * num_devices for step in range(1000): key, subkey jax.random.split(key) x, y make_batch(subkey, global_bs) # 关键 reshape第一维切成设备数pmap 会沿这一维映射 x x.reshape(num_devices, per_device_bs, -1) y y.reshape(num_devices, per_device_bs, -1) params train_step(params, (x, y)) if step % 100 0: # 取第一张卡的参数和对应 batch 计算 loss single_params jax.device_get(params)[0] loss loss_fn(single_params, jax.device_get(x)[0], jax.device_get(y)[0]) print(fstep {step}, loss{loss:.6f})这段代码里有几个极容易出错的地方我逐一说明第一train_step里的lax.pmean是必须的。它的作用是把所有设备上算出的梯度求平均然后广播回每张卡保证所有设备拿着完全相同的梯度去更新参数。如果你不写这行各设备会各训各的loss 可能还会下降但每个设备上的模型已经是不同的模型了。第二x.reshape(num_devices, per_device_bs, -1)的语义是“把原来一个大的 batch 切成 N 份每个设备拿其中一份”。这里不能用x.reshape(num_devices, -1, x.shape[-1])之外的写法乱改否则切片顺序会和设备对应关系错乱导致每个设备拿到的数据不均匀。第三多设备下打印 loss 不能直接取params整体。因为params是一个跨设备复制的分片对象需要先用jax.device_get把它搬回主机内存再取第一个设备的副本计算 loss。我刚开始直接loss_fn(params, x, y)会报 shape 不匹配的错最后才发现是忘了 device_get。第四jax.random.split(key)在训练循环里每次都要重新分裂出新 key。JAX 的随机数不像 PyTorch 那样有全局状态必须显式传递。特别是分布式的场景如果所有设备共用同一个 key不同设备的 batch 会一模一样数据并行就退化成“复制训练”梯度计算完全无效模型不会收敛。4.3 训练参数需要关注的细节从单卡切到多卡时最容易被忽视的是 global batch size。我在代码里用global_bs per_device_bs * num_devices定义了全局 batch然后把它切成设备维。这样每张卡的实际 batch size 依然是per_device_bs数学上等价于把一个大 batch 分到多张卡上算梯度。学习率的选择也值得讨论。如果你保持单卡时的学习率不变但把 global batch 扩大 N 倍经验上调优空间会变得很窄。常见的做法是“线性缩放规则”global batch 翻倍学习率也近似翻倍但实际项目里需要结合 warmup 来防止前期不稳定。我这里就没展开调优保持一个固定学习率只是为了验证分布式逻辑正确性。还有一个工程细节optax.apply_updates返回的是新的参数 pytree它和旧参数是完全不可变的两个对象。在分布式循环中我每次把返回值重新赋给params不需要像 PyTorch 那样zero_grad()、step()分开调。这也是为什么 JAX 训练循环能写得很紧凑的原因之一。4.4 多机多卡coordinator 模式上面的脚本在多机场景下只需要改两处初始化部分和启动方式。初始化部分我用jax.distributed.initialize()按RANK、WORLD_SIZE环境变量切分前面已经展示过。启动方式上建议用类似 PyTorch 的torchrun的思路用一个统一命令在多台机器上拉起脚本# 机器 0 RANK0 WORLD_SIZE2 COORDINATOR_ADDRESS10.0.0.1:1234 python train.py # 机器 1 RANK1 WORLD_SIZE2 COORDINATOR_ADDRESS10.0.0.1:1234 python train.pyJAX 底层的进程间通信同样依赖 NCCL所以多机网络带宽依然很重要数据加载的 IO 和网络通信的重叠策略也要设计好。JAX 的好处是通信模式主要由编译器决定你不需要手写send、recvNCCL 的调用是在 XLA 编译阶段自动生成的。5. 从 PyTorch 迁移到 JAX 的常见坑与排查技巧5.1 PRNG key随机数世界观的改变这是迁移过程中第一个冲击。PyTorch 的随机数基于全局状态torch.manual_seed(42)之后每次torch.randn都会按顺序产生新随机数非常符合直觉。JAX 的随机数则是显式状态机必须有一个PRNGKey对象每次生成随机数前用jax.random.split分裂出新 key然后用新 key 生成数据。在训练循环里我习惯在每个 step 内做key, subkey jax.random.split(key) x, y make_batch(subkey, global_bs)这保证了每个 step 使用不同的随机数。如果忘了 split同一个 step 里所有需要随机数的操作会拿到完全相同的随机数数据增强、dropout、数据采样都会异常。5.2 别忘了 pmeanpmap 不会自动做梯度平均这是 JAX 分布式训练与 PyTorch DDP 最大的行为差异。PyTorch DDP 在梯度同步时默认把所有 rank 的梯度求平均而 JAX 的pmap只是把你的函数复制到多个设备执行它不知道你算是要“平均”还是“求和”。你需要用lax.pmean或lax.psum明确告诉编译器你对哪个轴做归约。我在第一次跑多机实验时因为没有 pmean8 台机器各训各的loss 单卡看起来在降但验证集上模型几乎不动。排查了很久才发现是梯度没有平均值。从那以后我给自己定了个检查清单只要用了pmap第一步就看函数里有没有axis_name对应的归约操作。正确写法def train_step(params, batch): grads grad(loss_fn)(params, *batch) grads lax.pmean(grads, axis_namedevices) # 必须与 pmap 的 axis_name 一致 return optax.apply_updates(params, grads) train_step pmap(train_step, axis_namedevices)5.3 编译缓存与显存控制JAX 的首次jit/pmap编译时间可能很慢这是正常的因为 XLA 要生成针对具体设备的优化代码。一旦编译完成后续调用会命中缓存速度就上去了。但如果你的代码里每次传入的 batch shape 变化JAX 会认为计算图发生了变化而重新编译训练速度会断崖式下降。显存控制也是常见问题。XLA 默认会预分配大块显存来提升性能这在多卡并行时可能导致显存不足。你可以通过环境变量控制export XLA_PYTHON_CLIENT_MEM_FRACTION0.9或者改用平台分配器export XLA_PYTHON_CLIENT_ALLOCATORplatform我自己的经验是训练显存峰值不稳定时先调XLA_PYTHON_CLIENT_MEM_FRACTION0.95如果仍然 OOM再考虑减少 per_device batch size。单纯调低这个比例会影响性能需要权衡。5.4 一个速查表现象根本原因解决办法pmap报错 “sharded arguments” 相关batch 第一维没有切成设备数x.reshape(num_devices, per_device_bs, -1)多卡 loss 下降但模型不收敛忘了梯度归约各设备各自更新在 train_step 里加lax.pmean所有设备 batch 一模一样用了同一个 PRNG key 生成数据每次循环key, subkey jax.random.split(key)首次训练极慢JAX 在编译 XLA 图正常现象等编译结束保持 batch shape 固定避免重复编译显存溢出XLA 预分配策略XLA_PYTHON_CLIENT_MEM_FRACTION0.9或ALLOCATORplatform在jit函数里 print 看不到值函数被编译成计算图用jax.debug.print或在 eager 模式下调试换成 PyTorch 常用到的场景假如你在 PyTorch 环境遇到conda activate pytorch显示无法识别 conda 命令本质是 conda 所在的 PATH 没配置好这不属于模型训练问题但新手经常被困住。这类环境问题我觉得最好的解法是统一用 Python 虚拟环境或 Docker把 conda 的 PATH 烦恼降到最低。JAX 的环境抽象其实比 PyTorch 更依赖一个干净的 Python 环境所以我也会尽量把依赖装进同一个带锁文件的虚拟环境里避免项目之间互相拉扯版本。最后分享一个我个人的体会如果你之前主力是 PyTorch刚切到 JAX 的前几天一定要忍住“用熟悉的那套方式硬套”的冲动。JAX 函数式、不可变、显式随机数这套设计一开始确实会增加学习成本但等你能很自然地在脑子里把pmap想成“编译器层面的数据并行”再回头看 PyTorch 的 DDP会觉得两者其实是解决同一个问题的两条不同路线。我的建议也很简单先把一个小模型在 JAX 上完整跑通再逐步加jit、pmap、多机初始化每层都验证一次。按这个顺序走你会比直接拿一个大型分布式脚本硬啃快得多。