企业官网建设流程全解析

聊起分布式训练，很多人的第一反应是 PyTorch 的 DDP，进阶一点会碰到 FSDP。但如果你最近一年多在开源社区里刷模型代码，应该会注意到另一条主线——JAX。JAX 的分布式训练和 PyTorch 的差异，不是少了几个 API 那么简单，而是整套心智模型都不一样。在 PyTorch 里，你会觉得分布式是“挂”在训练循环上的：模型包一个 DDP，进程组初始化一下，梯度就自动同步了。在 JAX 里，分布式是计算图编译期就定下来的事情，你写的是一个“并行程序”，而不是“先写单机再想办法并行”。这篇文章我会从设计哲学讲到代码实现，给你一份可以直接对照的 JAX 分布式训练指南，也把 PyTorch 老手刚切换过来最容易踩的坑列清楚。不管你是正在做技术选型调研，还是已经被项目逼着迁移，这篇都值得先看完再动手。

1. 先框定讨论范围：分布式训练到底涉及哪些事

1.1 常见并行范式与它们各自的难点

聊 JAX 之前，得先把“分布式训练”这个词掰开。往简单了说，分布式训练就是三件事：数据怎么拆，模型怎么放，梯度怎么汇总。

数据并行最直观。每张卡拿同一个模型的完整副本，喂不同批次的数据，各自算梯度，然后做一次 AllReduce 把所有梯度加起来求平均，再同步更新参数。PyTorch 的 DDP 做的就是这件事，FSDP 则是把参数、梯度、优化器状态也分片存储的更进阶版本，目的是让单卡显存放得下超大模型。模型并行更像把模型横向切段，不同的层放在不同的设备上，数据按顺序在各段之间流转。张量并行更细，连一个矩阵乘法的行和列都要切成多份，同时交给多张卡计算，最后再拼起来。

这三种并行方式，难点其实不在“怎么切”，而在“切完之后的通信怎么组织”。什么时候发起通信、通信量和计算怎么重叠、不同并行方式叠加在一起时消息怎么路由，这些才是工程上真正复杂的地方。PyTorch 生态里，DDP 处理数据并行很成熟，FSDP 也封装好了不少逻辑，但模型并行和张量并行往往要借助 Megatron 这类外部代码或者手写 device mesh 来管理。JAX 想解决的问题，恰恰就是把这一大堆方案统一成一两套抽象，让“并行”不再是事后打补丁的工程代码，而是编译期就能理解、能优化的一等公民。

1.2 JAX 的核心哲学：把并行方案交给编译器

JAX 底层是 XLA，XLA 的思路是先把整个计算过程编译成一个完整计算图，再做全局优化。叠加到分布式场景时，XLA 甚至把跨设备的集合通信也当作图里的普通算子来对待。你写一个lax.pmean，编译器知道这是一个需要在设备间做全局求平均的操作；你写一个分片标注，编译器知道这个张量被拆到了哪些设备上。最后它会结合设备编号、通信组的拓扑位置，生成真实可执行的 NCCL 或 MPI 调用。

这个抽象层级带来的直接结果，是 JAX 的训练代码里很少出现“显式同步梯度”这种句子。你在 PyTorch 里调用loss.backward()之后，DDP 会在背后安排梯度 AllReduce；而 JAX 里，你在损失函数和参数更新逻辑里声明“梯度需要做 pmean”，剩下的事情由编译期规划。从这个意义上说，PyTorch 的分布式是运行时的艺术，JAX 的分布式是编译期的科学。我后来发现，这个差异几乎决定了你后续所有调试和优化手段的选择。

2. JAX 和 PyTorch 最大的思维差异在哪

2.1 PyTorch 的动态图心智：对象、运行时和隐式状态

PyTorch 默认是动态图，代码一行一行执行，执行到loss.backward()时，框架沿着 autograd 图反向传播。模型本身是nn.Module的实例，参数以.parameters()的形式挂在对象上。DDP 做的事情，本质上是把这个对象复制到每张卡上，在反向传播过程中插入梯度同步的 hook。对使用者来说，改造流程很固定：先初始化进程组，再用DDP包一层模型，数据搬到.cuda()，剩下语法与写单机几乎一致。

这个设计最大的优点是容易上手、容易调试。你可以打印中间张量，可以随时打断，甚至可以在反向传播中途插入自定义逻辑。但缺点也随之而来：框架在运行前看不到完整计算图，大量通信调度只能发生在运行时，由一个个 kernel 启动之间的顺序保证。模型规模变大之后，通信与计算的重叠优化、多种并行方式的自动规划，框架能帮忙的部分就很有限了。PyTorch 社区通常通过 FSDP、tensor parallel 库、或者干脆手写通信原语来弥补，但这也意味着上层要做更多补丁。

2.2 JAX 的函数式建模：把一切状态都摊开在明面上

JAX 的核心假设是：整个训练过程就是一个纯函数。你输入一组参数和一批数据，它返回新参数和一个 loss。模型没有“对象内部状态”，没有self，不鼓励在训练中保存任何可变全局变量。参数通常是 pytree——可以是任意嵌套的 dict、list、tuple 结构里的叶子节点；优化器状态也是显式传入传出的。

从 PyTorch 刚切过来的朋友通常会很不适应：写惯了model.eval()、optimizer.step()这种改变对象状态的魔法，JAX 里压根没有这些东西。BatchNorm 的 running mean、Dropout 的随机状态，都得自己作为参数传，或者用显式的 state 机制管理。一开始确实繁琐，但训练分布式模型时，“纯函数 + 不可变状态”是巨大的优势。每张卡拿到的是同一个函数，你只需要告诉编译器数据怎么分，它就能推导出哪些通信需要发生。不会出现“某个 buffer 忘同步了”这种隐蔽 bug，因为一切都在数据流里摊开着。

2.3 同样是梯度同步，机制和时间点完全不同

如果把一次数据并行训练拆开，会发现两边的“梯度同步”从里到外都不一样。PyTorch DDP 在反向传播过程中，autograd hook 捕捉每个参数的梯度，把它按 bucket 分组，桶满或反传结束时启动一次 AllReduce。为了加速，DDP 还会把通信和反向计算重叠起来，但这些都属于运行时调度。JAX 的pmap或jit加sharding，则在编译期就把“梯度需要跨设备求平均”这步写进图里。你的训练函数里写了jax.lax.pmean(grads, axis_name='batch')，编译器据此插入对应通信节点，并且清楚这条通信边在整个依赖图中的位置。

也因此，两边性能排查的思路不同。PyTorch 下多卡加速比不理想，你多半会去看通信桶大小、梯度累积的时机、NCCL 的带宽利用；JAX 下则要去看分片规格是否合理、Mesh 轴空间是否造成过度通信、编译是否被不必要的大型 kernel 拖慢。我认识不少朋友从 PyTorch 切到 JAX 后，第一个月都在用错误的思路调优，根本原因是脑子里的执行模型还没换过来。

3. 数据并行最小实现对比：DDP 和 pmap

3.1 PyTorch DDP 的训练骨架

先看 PyTorch 最典型的一套 DDP 代码骨架：

import torch import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP def train(rank, world_size): dist.init_process_group(backend="nccl", rank=rank, world_size=world_size) torch.cuda.set_device(rank) model = MyModel().cuda(rank) ddp_model = DDP(model, device_ids=[rank]) optimizer = torch.optim.Adam(ddp_model.parameters(), lr=1e-3) for batch in dataloader: x, y = [t.cuda(rank) for t in batch] optimizer.zero_grad() loss = ddp_model(x, y) loss.backward() # DDP 在这里自动同步梯度 optimizer.step()

要点是三件事：用torchrun启动多进程，每个进程绑定一张卡；模型包一层 DDP 之后，参数结构对外保持不变；梯度同步发生在loss.backward()过程中，你不用手动触发。DDP 内部会做初始权重广播、梯度分桶、AllReduce 插入。真要调优时，static_graph、bucket_cap_mb、gradient_as_bucket_view这些参数就开始冒头了。不过大部分项目其实用不到它们，默认配置也能跑得不错。

PyTorch 的 DDP 也有自己的一些别扭点。最明显的是多机训练时进程管理比较繁琐，MASTER_ADDR、MASTER_PORT、RANK、WORLD_SIZE这些环境变量一旦没对齐，就会出各种“address already in use”或者“connection reset”之类的报错。调试起来，往往要在好几台机器上同时看日志。

3.2 JAX pmap 的训练骨架

早期 JAX 最经典的写法是pmap。它把函数编译成 SPMD 程序，在每个设备上复制一份执行，并允许用axis_name做集合通信。一个最小数据并行例子：

import jax import jax.numpy as jnp from jax import pmap, grad def compute_loss(params, x, y): pred = model_forward(params, x) loss = jnp.mean((pred - y) ** 2) return loss grad_fn = grad(compute_loss) @pmap(axis_name="batch") def train_step(params, x, y): grads = grad_fn(params, x, y) grads = jax.lax.pmean(grads, axis_name="batch") # 跨设备梯度求平均 new_params = jax.tree.map(lambda p, g: p - 0.001 * g, params, grads) return new_params, loss

pmap会自动把输入按第一个维度切分，分发到多张卡，也就是数据并行里的“切 batch”。pmean是 JAX 的集合通信原语，意思是把每个设备上计算出的grads做全局求和再平均，编译到时对应一次 AllReduce。这里没有“哪个进程是 master”的概念，每台设备执行的是同一个完整函数，主进程和从进程的划分被抽象掉了。

这段代码虽然能跑，但它有一个隐藏前提：每个设备都要有完整模型参数。如果参数超过单卡显存，pmap就无能为力了。这也是 JAX 官方后来逐步不再主推pmap的核心原因。

3.3 从两段代码里能直接看出的六个差异

把两段代码并排看，差异是很清晰的。

一个是“对象方法调用式”，一个是“函数式数据流”。PyTorch 让你习惯性地调用model(x)，JAX 则要求你显式地把params传进函数。状态管理上，PyTorch 参数藏在nn.Module里，JAX 参数以 pytree 形式显式进出。分布式控制的粒度也不同：PyTorch 用 DDP 包装器加 hook，JAX 用pmap和axis_name这样的函数级变换。调试体验更是两回事：PyTorch 可以一行行打断点，JAX 一旦被@pmap或@jit编译，中间过程基本是黑盒，报错信息也更偏底层。性能分析方面，PyTorch 可以直接用 profiler 观察 kernel 和通信事件，JAX 则要看编译日志、sharding 传播决策这些偏编译器的信息。最后是上手门槛，PyTorch 平缓很多，JAX 需要你先接受函数式思维。

下面这张表可以拿来快速对照：

4. 现代 JAX 的分片方式：sharding 取代 pmap

4.1 为什么 pmap 不够用了

pmap 有很好的教学价值，但它太偏数据并行。一开始用没什么问题，一旦模型参数超过单卡显存，比如几十 B 的大模型，参数本身要切分、梯度要切分、优化器状态也要切分，pmap 的表达力就不够了。JAX 团队也承认，pmap 对复杂张量分片的支持有限，更推荐用jax.sharding定义张量如何分布在设备上。

新写法的三项核心组件是：Mesh描述设备拓扑，PartitionSpec描述每个张量轴对应到哪个 mesh 轴，jax.jit的in_shardings和out_shardings参数告诉编译器输入输出应该怎么分布。分片之后的计算，XLA 会自动编译成带通信节点的分布式计算。这套组合比 pmap 更灵活，也更接近大模型训练里需要的多轴并行。

4.2 用 jax.sharding 做数据并行与模型并行

一个基于jax.sharding的数据并行写法大概是这样：

import jax import jax.numpy as jnp from jax.sharding import Mesh, PartitionSpec, NamedSharding devices = jax.devices() mesh = Mesh(devices.reshape((num_devices,)), ("dev",)) def put_in_device(x): spec = NamedSharding(mesh, PartitionSpec("dev")) return jax.device_put(x, spec) @jax.jit(in_shardings=put_in_device(x).sharding, out_shardings=...) def train_step(params, x, y): ...

实际项目里，你会给每个参数配置PartitionSpec。比如一个形状是[num_layers, hidden_dim]的矩阵，你希望在 8 卡上做张量并行，就把第一个轴切到"dev"上：PartitionSpec("dev", None)，意思是第一个维度按设备切分，第二个维度保持完整。编译器会根据这个标注自动决定中间张量和输出的分片方式，最后生成带通信节点的 XLA 图。

这套声明式写法的好处是，数据并行、模型并行、张量并行，乃至任意组合，都可以用同一套逻辑表达。你不用写进程组，也不用手动搬运张量。但代价是理解门槛更高。你一旦写错PartitionSpec，报错通常是在编译期抛出来的一堆 XLA 内部术语，而且很容易把问题掩盖在抽象层之下。

4.3 和 PyTorch FSDP 的对应关系

PyTorch 用户经常问我，JAX 里有没有类似 FSDP 的东西。答案是：有，但通常是你自己写出来的。PyTorch 的 FSDP 把参数分片、按需 all-gather、梯度 reduce-scatter 这些逻辑封装在模块内部，你换一层包装就得到分片效果。JAX 的理念则是，这些行为本质上都可以由Mesh加PartitionSpec表达。你不需要一个叫 FSDP 的库，只需要把参数按某种模式分片标注出来，编译期就会生成对应的通信步骤。

举一个简单的例子：二维 mesh，一根轴叫data，一根轴叫model。参数矩阵按PartitionSpec("model", None)切分，就实现模型并行；训练数据按PartitionSpec("data")切分，就实现数据并行。两者叠加之后，编译器会在训练过程中同时做梯度 AllReduce 和参数的 all-gather。这就是一个简化版 FSDP 的行为。JAX 给你很高的自由度，同时也要你自己理解和维护这些抽象。PyTorch FSDP 开箱即用，但你很难从外部看到内部每一步在做什么。这算是两种框架对“复杂”的不同取舍。

5. 实际操作中高频出现的坑与排查经验

5.1 编译期黑盒怎么应对

第一次用 JAX 做分布式，最容易被吓到的是编译耗时。模型不大的时候，@pmap或@jit首轮编译都要几十秒甚至几分钟。你会觉得“怎么这么慢”，但 XLA 确实在认真做全图优化，包括通信调度。训练进入稳态之后，每一步的开销主要在执行 kernel，调用损耗很低。我的建议是：先用小模型和小数据跑通分布式逻辑，确认分片和通信配置正确，再上大模型。不要一上来就在完整模型上试，否则编译五分钟、报错三行，你根本无从定位。

另外，JAX 的报错风格和 PyTorch 很不一样。它经常指向抽象语义层，比如 “axis_name not in ...”“partition spec mismatch”，不会直接告诉你“哪个参数没放对位置”。排查时，把jax.debug.visualize_array_sharding这类可视化工具有意识用起来，比反复读报错有效得多。或者干脆多打印输入输出的分片规格，先把数据分布确认了，再往下挖。

这里还要提醒一点：@jit编译后的函数内部中间张量并不会按 Python 语句逐字执行。你 print 到的，大概率只是传入编译前的数据，而不是计算过程中某个算子的输出。想在函数内部做print调试，要使用jax.debug.print这类专门原语，否则很容易得出错误结论。

5.2 数据加载和设备放置，看起来小实则很磨人

JAX 本质上把数据看作 NumPy 数组，不是 PyTorch 的 Tensor。虽然 JAX 支持 GPU 上的jnp.ndarray，但很多数据管线组件默认在 CPU 上。PyTorch 里你随手.cuda()一下就行；切到 JAX 后，要么用jax.device_put手动把数据放到对应的Sharding，要么用tf.data、PyTorch DataLoader把数据预处理好再搬到 JAX 设备。很多分布式训练部署到最后，瓶颈反而在 CPU 到 GPU 的数据搬运，而不是 AllReduce 通信。

我的一个实操经验是，先把数据管线单独搭好，确保输入 batch 和设备数对得上。比如你 8 卡并行，一次喂进去 64 条样本，JAX 按每设备 8 条切分。如果 batch 无法被设备数整除，或者最后一块数据不满一个 batch，就容易出现“call to pjit / custom partitioning”这类误导性报错。表面上问题在分片，实际根源在数据切分不均匀。

5.3 常见问题速查表

把我和身边同事在迁移过程中遇到的高频问题整理成了这个表，照着查能省不少时间：

5.4 性能调优的几条具体心得

性能调优，我第一建议永远是减少通信次数。数据并行场景下，梯度通信每一步都在发生，如果通信开销占了大头，可以试试增大 batch、减少同步步数，或者用梯度累积。第二步是合理设计 Mesh。节点内走 NVLink，节点间走普通高速网络时，把通信量大的轴尽量放在节点内，跨节点的轴只承担必要的大粒度通信，这样能明显压掉跨节点带宽瓶颈。

在 JAX 里还有一个特殊的调优点：尽量利用jax.jit的static_argnums参数，把不参与数据流的配置对象标记为 static，避免因为配置变化触发重新编译。模型定义里也要注意，多用jax.numpy而不是原生 Python for 循环。否则 Python trace 阶段会把循环展开成超长计算图，编译时间和显存占用都会成倍上涨。这些细节看上去都很零碎，但对我个人来说，恰恰是决定 JAX 分布式训练能不能跑出理想加速比的关键。

6. 选型建议与切换阶段的心理准备

6.1 什么样的项目适合直接用 JAX

技术选型没有银弹。如果你主战场是快速验证想法，团队对 PyTorch 生态已经熟得不能再熟，那继续 DDP/FSDP 没有任何问题。JAX 的优势主要体现在三个场景：超大模型需要精细控制通信和分片规划；项目追求极致编译优化，比如长期做推理性能挖掘；团队能够接受函数式编程风格，希望在状态管理、实验重放、分布式逻辑可组合性上获得更好的代码表达力。

如果团队完全没有 JAX 基础，却要在一个需要快速交付的项目里切换框架，前置学习成本就不只是写代码本身，而是整套编译执行模型的建立。我见过最快上手的，往往是平时就习惯写纯函数、会主动拆解状态的工程师；最难受的，是重度依赖nn.Module内部状态和 PyTorch 生态库的人。这不是贬低谁，只是两种思维取向的适配问题。

6.2 我自己切换时最深刻的几个体会

我自己的真实体会是，切换成本里最难的不是 API，而是写代码时脑中的执行模型。PyTorch 让你觉得程序是一行一行在跑的，JAX 让你把程序当成一次变换、一次编译。只要过了这道认知关，JAX 的分布式训练写起来反而更干净，因为它把 PyTorch 里依靠隐式 hook 和进程组约定完成的事情，直接摊开在你面前。摊开的过程一开始让很多人不舒服，但多跑几个项目之后，你会越来越喜欢这种确定感。

最后分享一个小技巧：想评估一套分布式改造方案可行性，别一上来就在完整模型上跑。先用jax.debug.visualize_array_sharding打印几个关键张量的分片，确认预期中的 Mesh 轴真的生效了，再让完整训练跑起来。这个“先可视化再全量跑”的习惯，是我从几次深夜 debug 里总结出来的。JAX 的分布式训练本身不难写，难的是你永远不要凭空想象分片逻辑是合理的，你需要让编译器告诉你它到底是怎么安排的。