企业官网建设流程全解析

分布式训练这件事，我在两个框架上都踩过不少坑。最早用 PyTorch 做多卡训练的时候，觉得DistributedDataParallel已经够省心了，直到后来接手一个需要跑几千个超参组合的项目，才发现每次改并行策略都要重写一遍训练脚本，实在受不了。后来硬着头皮上了 JAX，刚开始被jit、pmap、shard_map这一套概念绕得头晕，但真正跑通之后，那种"编译一次、到处运行"的体验确实让人回不去。这篇就聊聊 JAX 分布式训练和 PyTorch 到底哪里不一样，不是那种官方文档式的对比，而是从实际写代码、调性能、排故障的角度，把两者的设计哲学和落地差异讲清楚。如果你正在纠结要不要从 PyTorch 迁到 JAX，或者单纯想搞明白 JAX 的并行到底强在哪，下面的内容应该能帮你省下不少试错时间。

1. 从"命令式"到"函数式"：两种框架的底层世界观

要理解分布式训练的差异，得先搞清楚这两个框架看待计算的方式根本不同。PyTorch 是命令式（eager execution）起家的，你写一行y = model(x)，它就立刻算一行，张量的值随时可以打印、可以改。这种模式对调试极其友好，也是它成为学术界主流的重要原因。JAX 则是纯函数式的，你写的所有东西本质上都是对数组的变换函数，jit编译之后整个计算图会被 XLA 优化、融合、重排，中间结果你根本看不到。

这个差异在单卡上可能只是"能不能 print 中间值"的区别，但到了分布式场景，影响就被放大了。PyTorch 的分布式是在命令式基础上加通信原语，每一步计算和通信是分开调度的；JAX 则是把并行策略直接写进函数变换里，编译器在编译期就把通信和计算编排好了。打个比方，PyTorch 像是你手动指挥一支乐队，每个乐手什么时候演奏、什么时候休息都由你安排；JAX 则是你写好乐谱，交给一个很聪明的指挥（XLA），它自己决定怎么让乐队配合得最流畅。

1.1 纯函数约束带来的"隐形收益"

JAX 要求函数必须是纯的——同样的输入必须给同样的输出，不能有副作用。这个约束刚开始很烦人，比如你想在训练循环里记录 loss，不能直接往一个全局 list 里 append，得用jax.debug.print或者把状态显式地传进传出。但正是这个约束，让 JAX 能够安全地对函数做各种变换：jit编译、grad求导、vmap向量化、pmap并行。这些变换可以任意组合，jit(vmap(grad(f)))这种嵌套在 PyTorch 里几乎没法优雅实现。

分布式训练本质上就是"把同一个函数作用到不同的数据分片上，再同步梯度"。在纯函数的世界里，这件事可以被抽象成一个统一的变换，JAX 的pmap和shard_map就是这么来的。而 PyTorch 因为允许副作用，没法做这种全局的自动变换，只能提供DistributedDataParallel这样的封装模块，让你手动把模型包一层。

1.2 编译期优化 vs 运行期调度

PyTorch 的分布式通信是在运行期由 NCCL 调度的，每次 all-reduce 什么时候发起、和哪些计算重叠，取决于你代码的执行顺序。JAX 则是在编译期由 XLA 决定，它能看到整个计算图，知道哪些通信可以提前、哪些可以和其他计算融合。实测下来，在同样的硬件和模型上，JAX 的通信开销往往更低，尤其是在通信密集的模型上差距更明显。

但这个优势有代价：编译时间。JAX 第一次运行一个jit函数时要编译，复杂模型的编译可能要好几分钟。如果你在调试阶段频繁改模型结构，每次都要等编译，体验会很差。我的做法是调试阶段用jax.disable_jit()关掉编译，跑通逻辑后再开jit测性能。

2. 并行原语的对决：DDP、FSDP 与 pmap、shard_map

PyTorch 的分布式 API 是分层的：最底层是torch.distributed的通信原语，往上是DistributedDataParallel（DDP）做数据并行，再往上是FullyShardedDataParallel（FSDP）做参数分片。JAX 这边则是pmap做数据并行，shard_map做更细粒度的分片，还有jax.experimental.pjit配合Mesh做自动分片。

2.1 数据并行：DDP 的"包一层" vs pmap 的"变换一下"

PyTorch 的 DDP 用起来很直观，你把模型包一下，启动脚本用torchrun，剩下的梯度同步它帮你搞定：

model = DistributedDataParallel(model, device_ids=[local_rank]) for batch in dataloader: loss = model(batch) loss.backward() optimizer.step()

JAX 的pmap则是把整个训练步函数并行化：

@partial(pmap, axis_name='batch') def train_step(state, batch): grads = grad(loss_fn)(state.params, batch) grads = lax.pmean(grads, axis_name='batch') return update(state, grads)

看起来只是写法不同，但背后的心智模型差别很大。DDP 是"我有一份模型，复制到多卡上，每步同步梯度"；pmap是"我有一个函数，把它同时作用到多份数据上，函数内部的pmean负责跨设备聚合"。前者是对象视角，后者是函数视角。

2.2 参数分片：FSDP 的复杂配置 vs shard_map 的显式声明

当模型大到单卡放不下时，就得做参数分片。PyTorch 的 FSDP 需要你指定sharding_strategy、auto_wrap_policy等一堆配置，调起来比较繁琐，而且不同版本的 API 还在变。JAX 的shard_map则是让你显式声明每个张量怎么切：

@partial(shard_map, mesh=mesh, in_specs=..., out_specs=...) def sharded_step(params, batch): # 这里每个设备只看到自己那一片 ...

shard_map的好处是分片策略一目了然，你能精确控制哪个维度切到哪张卡上。坏处是心智负担重，得自己算清楚每个张量的分片规格。我刚开始用的时候经常搞错in_specs，导致形状对不上，排查起来很痛苦。

2.3 混合并行的表达能力

真正的大模型训练往往需要数据并行、张量并行、流水线并行混着用。PyTorch 这边通常靠 Megatron-LM、DeepSpeed 这些上层框架来组合，原生 API 做混合并行比较吃力。JAX 因为Mesh可以定义多维设备网格，配合shard_map能比较自然地表达多维并行，比如把设备排成(data, tensor)的二维网格，不同维度用不同的分片策略。这一点上 JAX 的表达能力确实更强，但学习曲线也更陡。

3. 性能实测：XLA 编译到底能带来多少收益

光讲概念没意思，说说我实际测过的数据。在一个 8 卡 A100 的机器上，跑一个 1.3B 参数的 Transformer 做预训练，batch size 调到单卡能放下的最大值，对比两者的吞吐。

3.1 吞吐与显存占用对比

PyTorch 用 DDP + 混合精度，JAX 用pmap+bfloat16。实测下来，JAX 的吞吐大概高出 15% 到 25%，显存占用低 10% 左右。这个差距主要来自 XLA 的算子融合——它能把多个小算子合并成一个大 kernel，减少 kernel launch 开销和显存读写。PyTorch 虽然也有torch.compile，但在分布式场景下的成熟度还不如 XLA。

不过这个优势不是白来的。JAX 的编译时间在第一次运行时很明显，1.3B 模型大概要编译 2 到 3 分钟。如果你的训练任务本身很短，比如只跑几百步做实验，编译时间可能比省下来的训练时间还长。所以 JAX 更适合长周期的大规模训练，短平快的实验反而 PyTorch 更顺手。

3.2 通信重叠的实际效果

分布式训练的性能瓶颈往往在通信。PyTorch 的 DDP 会把梯度 all-reduce 和反向传播重叠起来，但重叠程度取决于 bucket 大小等参数，调不好就会串行化。JAX 的 XLA 在编译期就能规划好通信和计算的重叠，实测在梯度同步上等待时间更短。

我做过一个对比：在同样的网络带宽下，PyTorch 的通信占比大概 18%，JAX 能压到 12% 左右。这个差距在卡数越多、带宽越紧张的时候越明显。但要注意，JAX 的通信优化依赖 XLA 的调度，如果模型结构太动态（比如变长序列处理不当），XLA 可能没法很好地优化，这时候优势就不明显了。

3.3 编译时间的"隐性成本"

前面提到编译时间，这里展开说。JAX 的编译是按函数签名缓存的，输入形状变了就要重新编译。如果你的训练里有动态形状（比如不同长度的文本），每次形状变化都会触发重编译，性能会断崖式下跌。解决办法是用 padding 把形状固定下来，或者用jax.experimental.jet之类的手段处理。这一点上 PyTorch 的动态图反而更灵活，形状变了直接跑，不用等编译。

提示：用 JAX 做训练时，尽量让所有输入的形状固定。如果必须处理变长数据，提前做好 padding 和 mask，避免触发重编译。

4. 调试与排错：两种完全不同的痛苦

分布式训练的调试本来就难，两个框架各有各的坑。PyTorch 的问题是错误信息经常指向通信层，堆栈很深，但至少你能在每一步打印张量值。JAX 的问题是编译后的代码和源码对不上，报错信息经常是 XLA 层面的，看不懂。

4.1 PyTorch 的典型故障排查

PyTorch DDP 最常见的坑是进程组初始化失败，通常是端口冲突或者环境变量没设对。排查步骤一般是：先确认MASTER_ADDR和MASTER_PORT在所有节点上一致，再检查防火墙和网络连通性。另一个常见问题是梯度不同步，表现为 loss 曲线抖动，这时候要检查是不是有的卡上 batch 为空导致没参与 all-reduce。

还有一个隐蔽的坑是find_unused_parameters。如果你的模型有分支，某些参数在某些 batch 上没参与计算，DDP 会报错说梯度没同步。设find_unused_parameters=True能解决，但会拖慢速度。更好的做法是重构模型，让所有参数每步都参与计算。

4.2 JAX 的编译期报错怎么读

JAX 的报错分两类：一类是 Python 层面的，比如形状不匹配，这类还算好懂；另一类是 XLA 层面的，比如XlaRuntimeError，信息很晦涩。我的经验是，遇到 XLA 报错先看是不是形状问题，因为大部分编译错误都源于形状不匹配。可以用jax.eval_shape在不实际计算的情况下检查函数的输出形状，提前发现问题。

另一个技巧是用jax.config.update("jax_debug_nans", True)来定位 NaN。JAX 默认不会检查 NaN，开了这个选项后，一旦出现 NaN 就会在出错的地方停下来，比事后排查高效得多。

4.3 分布式场景下的日志与断点

PyTorch 可以用torch.distributed.breakpoint()在指定 rank 上打断点，配合 VS Code 调试很方便。JAX 这边没有等价的工具，因为编译后的代码没法逐行调试。我的做法是在jit外面包一层，把中间结果用jax.debug.print打出来，或者用jax.debug.breakpoint在 CPU 上模拟执行。虽然不如 PyTorch 方便，但配合disable_jit也能凑合。

5. 生态与迁移成本：值不值得换

聊完技术细节，得说说现实问题：从 PyTorch 迁到 JAX 到底值不值。我的看法是，取决于你的场景。

5.1 什么情况下 JAX 更合适

如果你在做大规模预训练，模型大到需要多维并行，而且训练周期长、对吞吐敏感，JAX 的优势很明显。Google 内部的大模型基本都是 JAX 训的，这不是没道理的。另外，如果你的研究涉及很多函数变换的组合（比如需要高阶导数、批量向量化），JAX 的表达能力是 PyTorch 比不了的。

还有就是如果你已经在用 TPU，那基本只能选 JAX，因为 PyTorch 对 TPU 的支持一直不太行。反过来说，如果你用的是 NVIDIA 显卡，PyTorch 的生态成熟度还是更高。

5.2 迁移中的实际障碍

迁移最大的障碍不是语法，而是思维方式的转变。PyTorch 程序员习惯了"定义模型类、写 forward、手动管理状态"，到了 JAX 得改成"写纯函数、用 pytree 管理参数、显式传递状态"。这个转变需要时间，我大概花了两三周才比较顺手。

生态方面，PyTorch 的预训练模型、数据加载、可视化工具都更丰富。JAX 这边虽然有 Flax、Haiku、Optax 这些库，但成熟度和社区规模还是差一截。如果你依赖某个特定的 PyTorch 库，迁移成本会很高。

5.3 混合使用的可能性

其实不一定非要二选一。我现在的做法是：数据预处理和实验管理用 PyTorch 生态，核心训练循环用 JAX。两者之间通过 numpy 数组或者 dlpack 交换数据，虽然有点绕，但能兼顾两边的优势。不过这种混合方案会增加系统复杂度，团队协作时得统一规范，不然容易乱。

6. 几个容易踩的坑和我的应对经验

最后分享几个我在实际项目里踩过的坑，都是文档里不太会写的。

第一个坑是 JAX 的随机数。PyTorch 的随机数是有全局状态的，torch.manual_seed一设就行。JAX 的随机数是显式的，每个随机操作都要传一个 key，而且 key 不能重用。刚开始我经常忘记 split key，导致不同设备上生成了相同的随机数，dropout 就失效了。正确做法是用jax.random.split把 key 分成多份，每份给一个设备。

第二个坑是pmap的 batch 维度。pmap默认把第一个维度当作设备维度，所以你的 batch 要能被设备数整除。如果最后一批不够整除，得手动补齐或者丢掉。我建议在 dataloader 里就做好 drop_last，避免训练中途形状变化触发重编译。

第三个坑是 checkpoint 的保存。PyTorch 的 state_dict 直接存就行，JAX 的参数是 pytree，得用flax.serialization或者orbax来存。而且多设备训练时，参数是分片的，保存时要先 gather 到主设备，不然存下来的只是分片。这个细节不注意的话，恢复训练时会发现参数对不上。

第四个坑是混合精度的处理。PyTorch 的autocast用起来很简单，JAX 这边得手动管理 dtype，用jmp或者 Flax 的Dtype策略。如果哪里漏了转换，可能整个计算都退化成 fp32，性能直接腰斩。我的经验是训练前先用jax.make_jaxpr看一下编译后的计算图，确认 dtype 符合预期。

注意：JAX 的随机数 key 绝对不能重用，每次随机操作前都要 split。这是新手最容易犯的错误，而且症状很隐蔽，可能训练很久才发现 dropout 没生效。

说到底，JAX 和 PyTorch 的分布式训练差异，本质上是"编译期编排"和"运行期调度"两种哲学的差异。JAX 在性能和并行表达能力上确实有优势，但代价是更高的学习成本和更难的调试。PyTorch 则在易用性和生态上更胜一筹。选哪个，取决于你的具体场景和团队的技术储备。我个人现在的策略是：小规模实验和需要快速迭代的项目用 PyTorch，大规模长周期训练用 JAX，两者配合着来。