前段时间我把一个模拟图像分类模型从单卡扩展到多机多卡前后折腾了一圈最大的感触是JAX把分布式训练的门槛压得非常低。如果你已经会写JAX单卡训练那这篇文章就非常值得看完——我会分享这次实践中最手把手的代码方案、多机环境的最小配置、以及那些报错文档里根本搜不到但实战一定会撞上的坑。标题说“超轻松”不是广告话术而是JAX和其他框架的分布式方案体验差异实在太明显。大部分分布式框架的核心工作在于“管理”——管理进程组、管理通信后端、管理数据分配方式。而JAX的核心思路是“转换”你把训练函数写成一个纯函数JAX帮你把它变成分布在多个设备上的版本。pmap和jit就是两个最常见的转换工具配合准确的设备分片信息框架会自动在计算图中插入集合通信你甚至不需要显式调用任何all-reduce。这篇文章适合两类人一类是刚接触JAX、想直接上手分布式训练的初学者另一类是被其他框架的分布式世界搞烦了、想看看有没有更省心的方案的老手。我尽量把每一步为什么这么做也讲明白避免你只抄代码、不涨经验。1. 为什么说JAX分布式训练可以很轻松1.1 分布式训练的常规弯路先说大多数人接触“分布式训练”时容易掉进的坑。很多人以为分布式难在数学其实分布式训练难在系统细节设备发现、集合通信、分片策略、梯度同步、容错检查点。这些听起来都不难但串在一起任何一个环节出问题都会让整个集群卡死或者静默出错。其他常见框架的做法通常需要你先建立一套进程组明确谁是主节点、谁是工作节点然后给每个进程分配全局rank和local rank最后还要担心后端是NCCL还是GLOO通信库初始化失败怎么排查模型参数怎么广播梯度反向传播之后又要怎么聚合。这一整套流程如果写成教程能写好几万字而且环境一变步骤就要调整。JAX走的路完全不同。它从设计上就把“分布式”当成数组分片sharding问题来处理。你在单卡上写的计算函数本身就是纯函数JAX可以对它做各种转换grad是求导转换jit是编译转换pmap是把函数映射到多设备的转换。转换是自动的、可组合的而不是一堆需要你手动调用的运行时API。用一句话概括你不用去管理“分布式系统”你只需要关心“数据放在哪个设备、数组怎么切分”。这个模型上的差异是JAX分布式“轻”的根本原因。1.2 懒人第一法则分片传播Sharding PropagationJAX能让分布式简单很大程度上依赖于一个机制叫分片传播。理论解释是当你用NamedSharding把一个输入数组放到多个设备上时jit编译器会追踪这个数组参与的所有计算自动决定中间结果的物理分布并自动插入需要的集合通信操作比如all-gather、all-reduce、reshuffle等。如果你用过其他框架你会很震撼这意味着你并不需要显式地把模型切到每张卡上。你只需要把模型参数复制到每张卡或者分片到每张卡把数据batch按设备维度切开然后训练循环代码和单卡时几乎一模一样——其他所有细节都由编译器fallback处理。这个特性再配合JAX的纯函数式风格就会产生很舒服的体验。比如你只要把参数用device_put放到一个分片布局上然后jit编译训练函数JAX会在需要同步梯度的地方自动加上pmean或者psum。这就是“写起来像单卡跑起来却已经是分布式”的真实现感。不过这里也有个容易被忽略的细节JAX编译器并不会聪明到自动知道“你希望数据并行”你必须通过sharding信息告诉它哪些维度是要被分片的。如果你只是跑一个普通的单卡代码然后放到多卡上它大概率会把模型参数复制到所有设备然后让每个设备独立计算最后不通信——这是合法的但完全没生效。所以“轻松”不是没有条件的轻松的前提是你要理解分片标注。这个也好学我们下面直接上手。2. 数据并行实操从单卡到多卡的三种写法数据并行Data Parallelism是分布式训练里最常见、最好理解的一种并行模式每张卡上都有一份完整的模型参数每个设备处理不同的数据切片前向和反向独立计算最后梯度做一次全局平均再用平均后的梯度更新每张卡上的参数副本。我用这次模拟图像分类项目来演示。模型就是一个简单的两层MLP数据集就是常见的28x28图像分类任务。第一版是纯单卡代码逻辑非常简单import jax import jax.numpy as jnp import optax from jax import random def mlp(params, x): w1, b1, w2, b2 params h jnp.tanh(x w1 b1) y h w2 b2 return y def init_params(key): key1, key2 random.split(key, 2) w1 random.normal(key1, (784, 256)) * 0.01 b1 jnp.zeros(256) w2 random.normal(key2, (256, 10)) * 0.01 b2 jnp.zeros(10) return (w1, b1, w2, b2) def loss_fn(params, x, y): pred mlp(params, x) return jnp.mean(optax.softmax_cross_entropy_with_integer_labels(pred, y)) def train_step(params, opt_state, x, y): grads jax.grad(loss_fn)(params, x, y) updates, opt_state optimizer.update(grads, opt_state, params) params optax.apply_updates(params, updates) return params, opt_state这个train_step函数就是典型的单卡版本。它的输入输出都是pytree也就是由数组组成的嵌套结构这给后续的并行转换提供了非常好的基础。下面我们用三种不同的方式来把它变成多卡版本。2.1 pmap写法经典但值得理解的方案pmap是JAX里最经典的多设备转换方式。它会把函数广播到所有设备上然后把输入数组按第0维拆开每个设备处理一块。在你完全不知道底层通信细节的情况下它也能帮你同步梯度。先把train_step稍微改一下加入一步梯度全局平均from jax import lax def train_step_parallel(params, opt_state, x, y): grads jax.grad(loss_fn)(params, x, y) grads lax.pmean(grads, axis_namedata) updates, opt_state optimizer.update(grads, opt_state, params) params optax.apply_updates(params, updates) return params, opt_state p_train_step jax.pmap( train_step_parallel, axis_namedata, donate_argnums(0, 1), )这里面的关键就在这一行lax.pmean(grads, axis_namedata)。为什么一定要加这一行因为pmap会把整个函数复制到N个设备上执行每个设备各自算出一份梯度。如果不做通信每个设备拿到的梯度只是自己那份数据的梯度而不是全局平均后的梯度。一旦你跳过pmean参数更新就会变成“各练各的”模型永远无法收敛到一个统一状态。pmean做的事情本质上是先所有设备把自己的梯度累加起来再做一下平均然后把平均结果广播给所有设备。整个过程看着像黑魔法但你别忘了它是在PMAP的上下文中才有意义的——axis_name就是那个在pmap调用时指定的维度标记编译器靠它来识别应该沿哪个设备维度做集合通信。调用方式也相对简单。因为pmap会自动沿第0维切分数据所以送进去的数组必须有设备数量的倍数个batch。训练循环batch jnp.array(image_batch) # shape: (num_devices, batch_per_device, ...) losses p_train_step(params, opt_state, batch, labels)这里有一点要特别注意pmap之后的函数返回值会“多出一个设备维度”。如果你传的batch第一个维度是设备数那么返回的loss容器shape就是(num_devices, batch_per_device, ...)打印loss的时候要取loss[0]或者做一次平均。这个特性很容易让新手蒙圈。另外还有一个随机数的坑。如果你在训练函数内部生成了随机数比如dropout mask、数据增强的随机噪声你需要先用jax.random.split(key, num_devices)为每个设备单独生成一个key。否则所有设备会拿到完全相同的随机序列等于破坏了这个设备上的数据独立性。2.2 基于jit的现代写法只需要一行分片标注pmap的年代比较早了现在JAX里更推荐的是用jit配合显式的分片标注。这种方式可以让编译器自动处理通信细节而且不像pmap那样对函数结构和返回值有限制。如果你看过JAX最近几个版本的更新就会发现官方示例里pmap出现的频率越来越低越来越多的代码是jax.jit加NamedSharding来完成相同的分布式逻辑。我们要做的事情其实就三步定义一个分片布局NamedSharding把数据放到这个分片上在jit编译的训练函数里用jax.lax.pmean在指定axis上做归一化。直接看代码from functools import partial from jax.sharding import NamedSharding, PartitionSpec as P devices jax.devices() # 把所有设备组成一个一维mesh维度名称为 data mesh jax.sharding.Mesh(devices, (data,)) # 在 data 维度上分片意思是这个数组会被切成 len(devices) 份 data_sharding NamedSharding(mesh, P(data)) partial( jax.jit, in_shardings(data_sharding, data_sharding), out_shardings(data_sharding, data_sharding), ) def p_train_step_jit(params, opt_state, x, y): grads jax.grad(loss_fn)(params, x, y) grads lax.pmean(grads, axis_namedata) updates, opt_state optimizer.update(grads, opt_state, params) params optax.apply_updates(params, updates) return params, opt_state这里in_shardings告诉jit编译器第一个参数params和第二个参数opt_state都应该按照data_sharding这种分片方式放置。由于params是一个pytree分片标注会递归作用到每个叶子上。具体来说是每个叶子数组都会沿着第0维被均匀分到所有设备上。而x和y在传入之前其实也需要是分片状态。最简单的做法是直接用jax.device_put(batch, data_sharding)x_sharded jax.device_put(x_batch, data_sharding) y_sharded jax.device_put(y_batch, data_sharding)这样jit编译时就知道输入数据已经是分片的后续计算就能正确进行。你可能会问这个写法比pmap先进在哪最大的优势是jit可以自动处理分片。编译器会根据你给的输入分片结合计算图上的算子推导出哪些中间结果应该分片、哪些应该复制并在必要的时候插入通信原语。而pmap的分片规则几乎是写死的——一切都沿着第0维拆。在某些复杂场景比如多维mesh、卷积的通道维切分、混合并行策略pmap就不够灵活了但jit配合sharding能够表达更复杂的并行策略。还有一个实际体验上的好处jit出来的编译产物可以复用。pmap每次输入shape变化都可能重新编译而jit静态shape优化后的执行速度更稳定。实测下来这个版本在多机环境下比pmap版本更稳日志里也不容易出现奇怪的shape不匹配报错。2.3 实际上手时的布局选择我们一开始做主模型实验的时候总想着一上来把所有参数都切到多个设备上。后来发现对小模型来说参数复制到每张卡才是最经济的。这句话值得单独拎出来说明数据并行有两种常见的参数摆放方式。一种是每张卡都放一份完整参数副本梯度平均后各自更新另一种是把参数本身也切到不同设备上用时再取。前者称为复制式数据并行Data Parallelism简称DP后者称为分片式数据并行即FSDP基础形态。对于这个模拟图像分类模型来说参数总量只有几十万到几百万个浮点数完整复制到每张卡完全无压力。复制式数据并行的好处是通信量少每个step只需要同步梯度也就是每个参数对应一个浮点数所有设备的梯度求和再广播。这个通信量约占参数量的两倍对网络压力很小。如果你的模型已经到了几十亿参数单卡显存放不下完整模型那就必须考虑分片参数甚至同时做流水线并行。那种场景下每张卡只持有模型的一部分参数前向计算时通过all-gather把所需参数拿全计算完后立刻释放参数。这个方向叫FSDPJAX里也能实现不过我这次先按住不讲那么深。先把复制式数据并行用熟已经能解决大部分中小模型的训练需求。我个人的建议是如果你的模型单卡能放下就用复制式数据并行简单、稳定、调参方便如果你的模型单卡放不下再去考虑FSDP或者混合并行那是另外一种工程问题。3. 多机多卡与内存优化细节3.1 多机集群初始化单机多卡只需要jax.device_count()看到所有本地GPU就行但多机训练就必须要告诉每个进程集群里有哪些机器、自己是谁。JAX提供的方法非常轻量不像很多框架需要启动脚本和复杂的配置。假设你有两台机器IP分别是192.168.1.10和192.168.1.11这里就用文档常见的示例IP实际换成你自己的内网地址。第一台机器作为协调者coordinator所有进程在启动时这样初始化import jax jax.distributed.initialize( coordinator_address192.168.1.10:1234, process_id0, # 当前进程编号 process_count2, # 总进程数 )第二台机器上的进程改成process_id1即可。初始化完成后jax.device_count()返回的就是两台机器上的所有可见GPU总数。这个API真的让人舒服。不需要额外安装NCCL的python包不需要设置WORLD_SIZE和RANK环境变量也不需要担心进程组是否创建成功。JAX后台会在初始化时自动完成设备发现和集群注册。需要注意的坑是所有机器的JAX版本、CUDA版本、NCCL版本必须一致。如果版本不一致常见现象是初始化时某台机器一直卡在等待状态没有任何报错过几分钟才超时。我曾经在这个问题上浪费了一个下午最后发现是其中一台机器的NCCL小版本不同。另外coordinator_address的端口要提前确认没有被防火墙拦住而且只允许一个节点做协调者。多节点环境下不要让每个节点都启动自己的分布式初始化否则会互相等待。3.2 分布式初始化之后的数据和标签加载多机训练中数据的正确分配是另一个重要环节。JAX不会自动帮你从数据中心读取数据再决定分给谁你需要自己保证每个进程读到不同的数据块。我们的做法是每个进程用jax.process_index()获取自己的全局进程编号然后基于进程编号做数据切分。比如使用tf.data或者PyTorch的DataLoader就按process_count做shard然后每个shard内部再按设备数切batch。这样每个设备在每个step都看到互不重叠的数据子集。如果你偷懒所有进程都从同一个文件读取全部数据然后直接分到设备上那你在多机场景会遇到大量重复数据、训练loss震荡、验证集指标虚高的问题。这个问题很难定位因为它不像报错那样明显所以建议在一开始就把数据切分的框架搭好。3.3 梯度同步与梯度累积数据并行里每个step的通信量是固定的每个参数一个数值需要聚合。但是通信频率很高——每个微型batch算完梯度就要同步一次。如果网络带宽相对显存容量不够通信本身反而可能成为瓶颈。这时候梯度累积Gradient Accumulation就派上用场了。思路很简单连续计算N个微型batch的梯度先不更新参数把梯度在本地累加等累加够N个之后再统一做一次全局聚合和参数更新。相当于把N次通信合并成1次。JAX实现梯度累积也很自然。核心思想是在外层累积梯度pytree达到累积步数后再更新def accumulated_train_step(params, grad_accum, opt_state, x, y): grads jax.grad(loss_fn)(params, x, y) grad_accum jax.tree.map(lambda ga, g: ga g, grad_accum, grads) return grad_accum在循环里每N步调用一次优化器更新。这能有效降低通信频次尤其是在多机环境下。我建议把累积步数设成2或4起步配合jax.debug的profile工具看一下训练吞吐再按需调整。优化器状态在梯度累积过程中不能更新。很多新手会把梯度累积放在优化器更新之后那等于白累积——参数已经更新过了梯度再累积就乱了。正确顺序是先用纯函数累积梯度累积足够了再调optimizer.update。3.4 混合精度与显存优化多卡最直接的好处就是显存扩容。单卡放不下的模型切成4卡分片基本可以马上再大一个数量级。但在实际运行中我们经常先碰到的不是模型放不下而是batch size受限于单卡显存。解决batch size限制的常见手段是混合精度。JAX的dtype控制非常灵活可以用参数化的dtype精确控制模型中间结果的类型。简化的做法是模型前向计算使用jnp.float16或者jnp.bfloat16反向计算梯度时保持float16但优化器状态和主权重保留float32。这样既保证训练数值稳定性又能把激活值显存占用减半。JAX里的代码写法是这样def loss_fn_mixed(params, x, y): # 选一个参数模拟到fp16 half_params jax.tree.map(lambda p: p.astype(jnp.float16), params) pred mlp(half_params, x.astype(jnp.float16)) loss jnp.mean(optax.softmax_cross_entropy_with_integer_labels(pred, y)) return loss注意这里计算是fp16但优化器更新还是用fp32的params。一开始直接用原生混合精度就容易遇到loss爆炸的问题后来我习惯加上grad clip效果稳定很多。Fishy的地方在于loss的dtype也变成fp16了数值范围很容易溢出。实践中可以在损失函数末尾加一句loss loss.astype(jnp.float32)让梯度保持在fp32范围。4. 常见问题与排查技巧实录4.1 写在前面排错的第一原则分布式训练排错的最大敌人不是技术难而是信息不对称。你在一台机器上看到的错误可能和另一台机器上的完全不一样而且很多错误不会直接显示在stdout上而是某个进程静默退出、其他进程等待超时。所以在动手排查之前先做好两个动作第一统一日志级别。给每个进程加上进程编号前缀方便确定消息来自哪台机器哪个进程。第二先在小规模上复现问题。比如先用单机双卡跑一个极小的batch确认逻辑正确后再扩展到多机。这个习惯能帮你省下一半的排查时间。我自己踩过的坑很多都是“大集群上跑一个巨大模型报错后淹没在几千行日志里”。后来我学乖了任何新改动先缩到最小可复现例子上跑通再放大规模。这一条比下面任何具体技巧都重要。4.2 常见问题速查表我把这次项目中遇到的类型按特征整理成了表格直接对照排查错误特征可能原因解决方案多卡训练正常但正确率异常低数据未正确分片各设备看到重复数据检查进程索引按process_index做数据shard梯度更新后loss不下降使用了pmean但参数没有统一更新确认所有设备用平均后的梯度更新参数初始化阶段进程挂起coordinator地址配置错误或NCCL版本不一致检查地址、端口和所有机器的库版本报错“device count mismatch”各节点GPU数量不一致或设备不可见用nvidia-smi逐台确认保证所有进程可见设备数相同显存OOMbatch size过大或显式缓存没有被释放减小microbatch、启用梯度累积、或使用混合精度输出shape多了一维使用了pmap返回结果是按设备链接的在需要打印/保存时取第一个设备的结果或做reshape多机不能通信防火墙拦住了通讯端口统筹协调放行范围保证集群内部端口互通多机训练速度反而比单机慢通信频繁或数据加载变成瓶颈启用梯度累积、检查数据读取是否每个节点单独load、考虑更大batch模型参数更新不一致优化器状态未同步确保优化器状态也被pmean或带axis_name的聚合处理这张表是我在实操中总结的而不是从哪份文档里抄来的。每次遇到问题建议“先看现象再按表格里的可能原因逐条排除”效率比瞎猜高很多。4.3 日志与checkpoint的实际处理技巧训练循环里最常见的错误是每个step都打印所有设备上的loss。不要这么干。在一次多机训练中日志每秒刷几十行所有设备都在print结果就是看到各种时间交错的日志而且日志本身会拖慢训练速度。正确做法是周期性聚合一次loss比如每100个step做一次jax.numpy平均只在主进程打印。loss聚合本身也要注意通信成本。我的习惯是单独聚合一次scalar不要在每步都调用。在JAX里用jax.lax.all_gather或jax.lax.pmean聚合一个小数组即可。checkpoint的处理方式也和单卡不同。单卡可以直接保存params和opt_state。多卡模式下如果参数是复制式的每张卡完全相同你只需要保存一份即可直接取params[0]保存。如果参数是分片式的那就必须先做一次all_gather把所有分片收集到主进程再保存完整版本。保存时建议顺带保存process_index和device_count信息不然下次加载时不知道这份checkpoint是在多少卡环境下保存的。如果设备数变了加载checkpoint后需要重新做sharding布局适配。4.4 关于随机性的坑JAX的随机性是很严格的。传统框架里random.seed(0)能做到全局统一但JAX推荐用显式的PRNGKey每个设备用不同的子key。在数据并行训练里即便你已经用jax.random.split(key, num_devices)生成了每台设备独立的key如果不小心在分片之后又用了同一个key仍然会出现随机序列重复。最常见的隐藏坑是数据增强在设备上完成而device上每个rank的随机key没有正确分配。我的做法是在每个step开始前根据当前step的seed和process_index为每个设备生成独立的子key。这样即使多个进程启动的时间完全一致每台设备拿到的增强结果也不可能完全重复。结尾回到开头那句话JAX分布式训练确实轻松但这种轻松建立在“你愿意花一晚上理解分片而不是花一整天配置集群”的前提下。框架把复杂彻底藏起来了你要做的只有尊重它的设计逻辑把模型和优化器写成纯函数用分片信息表达你的并行意图剩下的事情——通信、同步、形状调整——让编译器去处理。这次做完之后我的体会是如果你想在项目里快速验证分布式训练JAX是我当前用过的所有工具里学习曲线最友好、迁移成本最低的选项。当然JAX也有它自己的生态成熟度和资料稀疏等问题但“易用性”这一项它确实做得很出色。如果你正打算把现有模型扩展到多卡不妨从今天分享的这个最小例子开始改跑通了就是胜利。 
长河君 · 2026/10/22 21:16:13 · 阅读 1.2k JAX分布式训练超轻松：从单卡到多机多卡的实战指南
前段时间我把一个模拟图像分类模型从单卡扩展到多机多卡前后折腾了一圈最大的感触是JAX把分布式训练的门槛压得非常低。如果你已经会写JAX单卡训练那这篇文章就非常值得看完——我会分享这次实践中最手把手的代码方案、多机环境的最小配置、以及那些报错文档里根本搜不到但实战一定会撞上的坑。标题说“超轻松”不是广告话术而是JAX和其他框架的分布式方案体验差异实在太明显。大部分分布式框架的核心工作在于“管理”——管理进程组、管理通信后端、管理数据分配方式。而JAX的核心思路是“转换”你把训练函数写成一个纯函数JAX帮你把它变成分布在多个设备上的版本。pmap和jit就是两个最常见的转换工具配合准确的设备分片信息框架会自动在计算图中插入集合通信你甚至不需要显式调用任何all-reduce。这篇文章适合两类人一类是刚接触JAX、想直接上手分布式训练的初学者另一类是被其他框架的分布式世界搞烦了、想看看有没有更省心的方案的老手。我尽量把每一步为什么这么做也讲明白避免你只抄代码、不涨经验。1. 为什么说JAX分布式训练可以很轻松1.1 分布式训练的常规弯路先说大多数人接触“分布式训练”时容易掉进的坑。很多人以为分布式难在数学其实分布式训练难在系统细节设备发现、集合通信、分片策略、梯度同步、容错检查点。这些听起来都不难但串在一起任何一个环节出问题都会让整个集群卡死或者静默出错。其他常见框架的做法通常需要你先建立一套进程组明确谁是主节点、谁是工作节点然后给每个进程分配全局rank和local rank最后还要担心后端是NCCL还是GLOO通信库初始化失败怎么排查模型参数怎么广播梯度反向传播之后又要怎么聚合。这一整套流程如果写成教程能写好几万字而且环境一变步骤就要调整。JAX走的路完全不同。它从设计上就把“分布式”当成数组分片sharding问题来处理。你在单卡上写的计算函数本身就是纯函数JAX可以对它做各种转换grad是求导转换jit是编译转换pmap是把函数映射到多设备的转换。转换是自动的、可组合的而不是一堆需要你手动调用的运行时API。用一句话概括你不用去管理“分布式系统”你只需要关心“数据放在哪个设备、数组怎么切分”。这个模型上的差异是JAX分布式“轻”的根本原因。1.2 懒人第一法则分片传播Sharding PropagationJAX能让分布式简单很大程度上依赖于一个机制叫分片传播。理论解释是当你用NamedSharding把一个输入数组放到多个设备上时jit编译器会追踪这个数组参与的所有计算自动决定中间结果的物理分布并自动插入需要的集合通信操作比如all-gather、all-reduce、reshuffle等。如果你用过其他框架你会很震撼这意味着你并不需要显式地把模型切到每张卡上。你只需要把模型参数复制到每张卡或者分片到每张卡把数据batch按设备维度切开然后训练循环代码和单卡时几乎一模一样——其他所有细节都由编译器fallback处理。这个特性再配合JAX的纯函数式风格就会产生很舒服的体验。比如你只要把参数用device_put放到一个分片布局上然后jit编译训练函数JAX会在需要同步梯度的地方自动加上pmean或者psum。这就是“写起来像单卡跑起来却已经是分布式”的真实现感。不过这里也有个容易被忽略的细节JAX编译器并不会聪明到自动知道“你希望数据并行”你必须通过sharding信息告诉它哪些维度是要被分片的。如果你只是跑一个普通的单卡代码然后放到多卡上它大概率会把模型参数复制到所有设备然后让每个设备独立计算最后不通信——这是合法的但完全没生效。所以“轻松”不是没有条件的轻松的前提是你要理解分片标注。这个也好学我们下面直接上手。2. 数据并行实操从单卡到多卡的三种写法数据并行Data Parallelism是分布式训练里最常见、最好理解的一种并行模式每张卡上都有一份完整的模型参数每个设备处理不同的数据切片前向和反向独立计算最后梯度做一次全局平均再用平均后的梯度更新每张卡上的参数副本。我用这次模拟图像分类项目来演示。模型就是一个简单的两层MLP数据集就是常见的28x28图像分类任务。第一版是纯单卡代码逻辑非常简单import jax import jax.numpy as jnp import optax from jax import random def mlp(params, x): w1, b1, w2, b2 params h jnp.tanh(x w1 b1) y h w2 b2 return y def init_params(key): key1, key2 random.split(key, 2) w1 random.normal(key1, (784, 256)) * 0.01 b1 jnp.zeros(256) w2 random.normal(key2, (256, 10)) * 0.01 b2 jnp.zeros(10) return (w1, b1, w2, b2) def loss_fn(params, x, y): pred mlp(params, x) return jnp.mean(optax.softmax_cross_entropy_with_integer_labels(pred, y)) def train_step(params, opt_state, x, y): grads jax.grad(loss_fn)(params, x, y) updates, opt_state optimizer.update(grads, opt_state, params) params optax.apply_updates(params, updates) return params, opt_state这个train_step函数就是典型的单卡版本。它的输入输出都是pytree也就是由数组组成的嵌套结构这给后续的并行转换提供了非常好的基础。下面我们用三种不同的方式来把它变成多卡版本。2.1 pmap写法经典但值得理解的方案pmap是JAX里最经典的多设备转换方式。它会把函数广播到所有设备上然后把输入数组按第0维拆开每个设备处理一块。在你完全不知道底层通信细节的情况下它也能帮你同步梯度。先把train_step稍微改一下加入一步梯度全局平均from jax import lax def train_step_parallel(params, opt_state, x, y): grads jax.grad(loss_fn)(params, x, y) grads lax.pmean(grads, axis_namedata) updates, opt_state optimizer.update(grads, opt_state, params) params optax.apply_updates(params, updates) return params, opt_state p_train_step jax.pmap( train_step_parallel, axis_namedata, donate_argnums(0, 1), )这里面的关键就在这一行lax.pmean(grads, axis_namedata)。为什么一定要加这一行因为pmap会把整个函数复制到N个设备上执行每个设备各自算出一份梯度。如果不做通信每个设备拿到的梯度只是自己那份数据的梯度而不是全局平均后的梯度。一旦你跳过pmean参数更新就会变成“各练各的”模型永远无法收敛到一个统一状态。pmean做的事情本质上是先所有设备把自己的梯度累加起来再做一下平均然后把平均结果广播给所有设备。整个过程看着像黑魔法但你别忘了它是在PMAP的上下文中才有意义的——axis_name就是那个在pmap调用时指定的维度标记编译器靠它来识别应该沿哪个设备维度做集合通信。调用方式也相对简单。因为pmap会自动沿第0维切分数据所以送进去的数组必须有设备数量的倍数个batch。训练循环batch jnp.array(image_batch) # shape: (num_devices, batch_per_device, ...) losses p_train_step(params, opt_state, batch, labels)这里有一点要特别注意pmap之后的函数返回值会“多出一个设备维度”。如果你传的batch第一个维度是设备数那么返回的loss容器shape就是(num_devices, batch_per_device, ...)打印loss的时候要取loss[0]或者做一次平均。这个特性很容易让新手蒙圈。另外还有一个随机数的坑。如果你在训练函数内部生成了随机数比如dropout mask、数据增强的随机噪声你需要先用jax.random.split(key, num_devices)为每个设备单独生成一个key。否则所有设备会拿到完全相同的随机序列等于破坏了这个设备上的数据独立性。2.2 基于jit的现代写法只需要一行分片标注pmap的年代比较早了现在JAX里更推荐的是用jit配合显式的分片标注。这种方式可以让编译器自动处理通信细节而且不像pmap那样对函数结构和返回值有限制。如果你看过JAX最近几个版本的更新就会发现官方示例里pmap出现的频率越来越低越来越多的代码是jax.jit加NamedSharding来完成相同的分布式逻辑。我们要做的事情其实就三步定义一个分片布局NamedSharding把数据放到这个分片上在jit编译的训练函数里用jax.lax.pmean在指定axis上做归一化。直接看代码from functools import partial from jax.sharding import NamedSharding, PartitionSpec as P devices jax.devices() # 把所有设备组成一个一维mesh维度名称为 data mesh jax.sharding.Mesh(devices, (data,)) # 在 data 维度上分片意思是这个数组会被切成 len(devices) 份 data_sharding NamedSharding(mesh, P(data)) partial( jax.jit, in_shardings(data_sharding, data_sharding), out_shardings(data_sharding, data_sharding), ) def p_train_step_jit(params, opt_state, x, y): grads jax.grad(loss_fn)(params, x, y) grads lax.pmean(grads, axis_namedata) updates, opt_state optimizer.update(grads, opt_state, params) params optax.apply_updates(params, updates) return params, opt_state这里in_shardings告诉jit编译器第一个参数params和第二个参数opt_state都应该按照data_sharding这种分片方式放置。由于params是一个pytree分片标注会递归作用到每个叶子上。具体来说是每个叶子数组都会沿着第0维被均匀分到所有设备上。而x和y在传入之前其实也需要是分片状态。最简单的做法是直接用jax.device_put(batch, data_sharding)x_sharded jax.device_put(x_batch, data_sharding) y_sharded jax.device_put(y_batch, data_sharding)这样jit编译时就知道输入数据已经是分片的后续计算就能正确进行。你可能会问这个写法比pmap先进在哪最大的优势是jit可以自动处理分片。编译器会根据你给的输入分片结合计算图上的算子推导出哪些中间结果应该分片、哪些应该复制并在必要的时候插入通信原语。而pmap的分片规则几乎是写死的——一切都沿着第0维拆。在某些复杂场景比如多维mesh、卷积的通道维切分、混合并行策略pmap就不够灵活了但jit配合sharding能够表达更复杂的并行策略。还有一个实际体验上的好处jit出来的编译产物可以复用。pmap每次输入shape变化都可能重新编译而jit静态shape优化后的执行速度更稳定。实测下来这个版本在多机环境下比pmap版本更稳日志里也不容易出现奇怪的shape不匹配报错。2.3 实际上手时的布局选择我们一开始做主模型实验的时候总想着一上来把所有参数都切到多个设备上。后来发现对小模型来说参数复制到每张卡才是最经济的。这句话值得单独拎出来说明数据并行有两种常见的参数摆放方式。一种是每张卡都放一份完整参数副本梯度平均后各自更新另一种是把参数本身也切到不同设备上用时再取。前者称为复制式数据并行Data Parallelism简称DP后者称为分片式数据并行即FSDP基础形态。对于这个模拟图像分类模型来说参数总量只有几十万到几百万个浮点数完整复制到每张卡完全无压力。复制式数据并行的好处是通信量少每个step只需要同步梯度也就是每个参数对应一个浮点数所有设备的梯度求和再广播。这个通信量约占参数量的两倍对网络压力很小。如果你的模型已经到了几十亿参数单卡显存放不下完整模型那就必须考虑分片参数甚至同时做流水线并行。那种场景下每张卡只持有模型的一部分参数前向计算时通过all-gather把所需参数拿全计算完后立刻释放参数。这个方向叫FSDPJAX里也能实现不过我这次先按住不讲那么深。先把复制式数据并行用熟已经能解决大部分中小模型的训练需求。我个人的建议是如果你的模型单卡能放下就用复制式数据并行简单、稳定、调参方便如果你的模型单卡放不下再去考虑FSDP或者混合并行那是另外一种工程问题。3. 多机多卡与内存优化细节3.1 多机集群初始化单机多卡只需要jax.device_count()看到所有本地GPU就行但多机训练就必须要告诉每个进程集群里有哪些机器、自己是谁。JAX提供的方法非常轻量不像很多框架需要启动脚本和复杂的配置。假设你有两台机器IP分别是192.168.1.10和192.168.1.11这里就用文档常见的示例IP实际换成你自己的内网地址。第一台机器作为协调者coordinator所有进程在启动时这样初始化import jax jax.distributed.initialize( coordinator_address192.168.1.10:1234, process_id0, # 当前进程编号 process_count2, # 总进程数 )第二台机器上的进程改成process_id1即可。初始化完成后jax.device_count()返回的就是两台机器上的所有可见GPU总数。这个API真的让人舒服。不需要额外安装NCCL的python包不需要设置WORLD_SIZE和RANK环境变量也不需要担心进程组是否创建成功。JAX后台会在初始化时自动完成设备发现和集群注册。需要注意的坑是所有机器的JAX版本、CUDA版本、NCCL版本必须一致。如果版本不一致常见现象是初始化时某台机器一直卡在等待状态没有任何报错过几分钟才超时。我曾经在这个问题上浪费了一个下午最后发现是其中一台机器的NCCL小版本不同。另外coordinator_address的端口要提前确认没有被防火墙拦住而且只允许一个节点做协调者。多节点环境下不要让每个节点都启动自己的分布式初始化否则会互相等待。3.2 分布式初始化之后的数据和标签加载多机训练中数据的正确分配是另一个重要环节。JAX不会自动帮你从数据中心读取数据再决定分给谁你需要自己保证每个进程读到不同的数据块。我们的做法是每个进程用jax.process_index()获取自己的全局进程编号然后基于进程编号做数据切分。比如使用tf.data或者PyTorch的DataLoader就按process_count做shard然后每个shard内部再按设备数切batch。这样每个设备在每个step都看到互不重叠的数据子集。如果你偷懒所有进程都从同一个文件读取全部数据然后直接分到设备上那你在多机场景会遇到大量重复数据、训练loss震荡、验证集指标虚高的问题。这个问题很难定位因为它不像报错那样明显所以建议在一开始就把数据切分的框架搭好。3.3 梯度同步与梯度累积数据并行里每个step的通信量是固定的每个参数一个数值需要聚合。但是通信频率很高——每个微型batch算完梯度就要同步一次。如果网络带宽相对显存容量不够通信本身反而可能成为瓶颈。这时候梯度累积Gradient Accumulation就派上用场了。思路很简单连续计算N个微型batch的梯度先不更新参数把梯度在本地累加等累加够N个之后再统一做一次全局聚合和参数更新。相当于把N次通信合并成1次。JAX实现梯度累积也很自然。核心思想是在外层累积梯度pytree达到累积步数后再更新def accumulated_train_step(params, grad_accum, opt_state, x, y): grads jax.grad(loss_fn)(params, x, y) grad_accum jax.tree.map(lambda ga, g: ga g, grad_accum, grads) return grad_accum在循环里每N步调用一次优化器更新。这能有效降低通信频次尤其是在多机环境下。我建议把累积步数设成2或4起步配合jax.debug的profile工具看一下训练吞吐再按需调整。优化器状态在梯度累积过程中不能更新。很多新手会把梯度累积放在优化器更新之后那等于白累积——参数已经更新过了梯度再累积就乱了。正确顺序是先用纯函数累积梯度累积足够了再调optimizer.update。3.4 混合精度与显存优化多卡最直接的好处就是显存扩容。单卡放不下的模型切成4卡分片基本可以马上再大一个数量级。但在实际运行中我们经常先碰到的不是模型放不下而是batch size受限于单卡显存。解决batch size限制的常见手段是混合精度。JAX的dtype控制非常灵活可以用参数化的dtype精确控制模型中间结果的类型。简化的做法是模型前向计算使用jnp.float16或者jnp.bfloat16反向计算梯度时保持float16但优化器状态和主权重保留float32。这样既保证训练数值稳定性又能把激活值显存占用减半。JAX里的代码写法是这样def loss_fn_mixed(params, x, y): # 选一个参数模拟到fp16 half_params jax.tree.map(lambda p: p.astype(jnp.float16), params) pred mlp(half_params, x.astype(jnp.float16)) loss jnp.mean(optax.softmax_cross_entropy_with_integer_labels(pred, y)) return loss注意这里计算是fp16但优化器更新还是用fp32的params。一开始直接用原生混合精度就容易遇到loss爆炸的问题后来我习惯加上grad clip效果稳定很多。Fishy的地方在于loss的dtype也变成fp16了数值范围很容易溢出。实践中可以在损失函数末尾加一句loss loss.astype(jnp.float32)让梯度保持在fp32范围。4. 常见问题与排查技巧实录4.1 写在前面排错的第一原则分布式训练排错的最大敌人不是技术难而是信息不对称。你在一台机器上看到的错误可能和另一台机器上的完全不一样而且很多错误不会直接显示在stdout上而是某个进程静默退出、其他进程等待超时。所以在动手排查之前先做好两个动作第一统一日志级别。给每个进程加上进程编号前缀方便确定消息来自哪台机器哪个进程。第二先在小规模上复现问题。比如先用单机双卡跑一个极小的batch确认逻辑正确后再扩展到多机。这个习惯能帮你省下一半的排查时间。我自己踩过的坑很多都是“大集群上跑一个巨大模型报错后淹没在几千行日志里”。后来我学乖了任何新改动先缩到最小可复现例子上跑通再放大规模。这一条比下面任何具体技巧都重要。4.2 常见问题速查表我把这次项目中遇到的类型按特征整理成了表格直接对照排查错误特征可能原因解决方案多卡训练正常但正确率异常低数据未正确分片各设备看到重复数据检查进程索引按process_index做数据shard梯度更新后loss不下降使用了pmean但参数没有统一更新确认所有设备用平均后的梯度更新参数初始化阶段进程挂起coordinator地址配置错误或NCCL版本不一致检查地址、端口和所有机器的库版本报错“device count mismatch”各节点GPU数量不一致或设备不可见用nvidia-smi逐台确认保证所有进程可见设备数相同显存OOMbatch size过大或显式缓存没有被释放减小microbatch、启用梯度累积、或使用混合精度输出shape多了一维使用了pmap返回结果是按设备链接的在需要打印/保存时取第一个设备的结果或做reshape多机不能通信防火墙拦住了通讯端口统筹协调放行范围保证集群内部端口互通多机训练速度反而比单机慢通信频繁或数据加载变成瓶颈启用梯度累积、检查数据读取是否每个节点单独load、考虑更大batch模型参数更新不一致优化器状态未同步确保优化器状态也被pmean或带axis_name的聚合处理这张表是我在实操中总结的而不是从哪份文档里抄来的。每次遇到问题建议“先看现象再按表格里的可能原因逐条排除”效率比瞎猜高很多。4.3 日志与checkpoint的实际处理技巧训练循环里最常见的错误是每个step都打印所有设备上的loss。不要这么干。在一次多机训练中日志每秒刷几十行所有设备都在print结果就是看到各种时间交错的日志而且日志本身会拖慢训练速度。正确做法是周期性聚合一次loss比如每100个step做一次jax.numpy平均只在主进程打印。loss聚合本身也要注意通信成本。我的习惯是单独聚合一次scalar不要在每步都调用。在JAX里用jax.lax.all_gather或jax.lax.pmean聚合一个小数组即可。checkpoint的处理方式也和单卡不同。单卡可以直接保存params和opt_state。多卡模式下如果参数是复制式的每张卡完全相同你只需要保存一份即可直接取params[0]保存。如果参数是分片式的那就必须先做一次all_gather把所有分片收集到主进程再保存完整版本。保存时建议顺带保存process_index和device_count信息不然下次加载时不知道这份checkpoint是在多少卡环境下保存的。如果设备数变了加载checkpoint后需要重新做sharding布局适配。4.4 关于随机性的坑JAX的随机性是很严格的。传统框架里random.seed(0)能做到全局统一但JAX推荐用显式的PRNGKey每个设备用不同的子key。在数据并行训练里即便你已经用jax.random.split(key, num_devices)生成了每台设备独立的key如果不小心在分片之后又用了同一个key仍然会出现随机序列重复。最常见的隐藏坑是数据增强在设备上完成而device上每个rank的随机key没有正确分配。我的做法是在每个step开始前根据当前step的seed和process_index为每个设备生成独立的子key。这样即使多个进程启动的时间完全一致每台设备拿到的增强结果也不可能完全重复。结尾回到开头那句话JAX分布式训练确实轻松但这种轻松建立在“你愿意花一晚上理解分片而不是花一整天配置集群”的前提下。框架把复杂彻底藏起来了你要做的只有尊重它的设计逻辑把模型和优化器写成纯函数用分片信息表达你的并行意图剩下的事情——通信、同步、形状调整——让编译器去处理。这次做完之后我的体会是如果你想在项目里快速验证分布式训练JAX是我当前用过的所有工具里学习曲线最友好、迁移成本最低的选项。当然JAX也有它自己的生态成熟度和资料稀疏等问题但“易用性”这一项它确实做得很出色。如果你正打算把现有模型扩展到多卡不妨从今天分享的这个最小例子开始改跑通了就是胜利。