4. 核心转换机制（下）：vmap向量化映射、pmap数据并行、grad自动微分、三者组合使用

好，咱们接着聊JAX的核心转换。上一节我们讲了jit和lax，这一节我把剩下的三个重头戏一口气讲完：vmap、pmap、grad。这三个东西，说白了就是JAX的「三板斧」——向量化、并行化、自动求导。我个人的经验是，把这仨玩明白了，你基本就能在硬件加速器上横着走了。

4.1 vmap：自动向量化，告别手写batch循环

先说说vmap。你想想看，写深度学习代码时，最烦的是什么？是手动写batch循环。每次都要把单个样本的处理逻辑，硬生生改成批处理版本。改完还得调试，调试完发现维度对不上，心态直接炸裂。

vmap就是来解决这个问题的。它能把一个只处理单个样本的函数，自动转换成能处理batch的函数。说白了，就是帮你自动加一个batch维度。

来看个例子。假设我们有一个简单的线性层函数：

import jax.numpy as jnp
from jax import vmap

def linear_layer(x, w, b):
    return jnp.dot(w, x) + b

# 单样本输入
x_single = jnp.ones((10,))
w = jnp.ones((5, 10))
b = jnp.ones((5,))

# 手动batch处理
def batch_linear(X, w, b):
    return jnp.stack([linear_layer(x, w, b) for x in X])

# 使用vmap自动向量化
batch_linear_vmap = vmap(linear_layer, in_axes=(0, None, None))

# 测试
X_batch = jnp.ones((32, 10))
result_manual = batch_linear(X_batch, w, b)
result_vmap = batch_linear_vmap(X_batch, w, b)

print(jnp.allclose(result_manual, result_vmap))  # True

注意看in_axes=(0, None, None)这个参数。它告诉vmap：第一个参数（x）沿着第0维映射，第二个和第三个参数（w和b）不映射，所有batch共享。我在项目中遇到过有人把in_axes全设成0，结果每个样本都用了不同的权重，训练直接崩了。嗯，这个坑我踩过。

4.2 pmap：数据并行，多设备一起上

vmap搞定了单设备上的向量化，那多设备呢？pmap就是干这个的。它把数据切分成多份，分发给多个设备并行计算。你想想看，8块GPU同时干活，速度理论上能快8倍。

pmap和vmap的区别在哪？vmap是在一个设备上串行计算batch中的每个样本（虽然内部有SIMD优化），pmap是把batch分到多个设备上并行计算。说白了，vmap是「一个人干多份活」，pmap是「多个人各干一份活」。

from jax import pmap
import jax

# 假设我们有4个设备
devices = jax.devices()[:4]

def compute_loss(params, batch):
    # 假设这是你的损失函数
    return jnp.mean(batch ** 2) + jnp.sum(params)

# 使用pmap进行数据并行
parallel_loss = pmap(compute_loss, in_axes=(None, 0))

# 准备数据：batch_size=32，切分成4份，每份8个样本
params = jnp.ones((10,))
batch = jnp.ones((32, 100))

# 注意：pmap要求batch维度能被设备数整除
result = parallel_loss(params, batch)
print(result.shape)  # (4,) 每个设备返回一个标量损失

pmap返回的结果，默认会在所有设备上各有一份。如果你需要聚合结果，可以用jax.lax.psum或者jax.lax.pmean做跨设备规约。我习惯在pmap的函数内部就做好聚合，这样外部拿到的就是最终结果。

4.3 grad：自动微分，反向传播的瑞士军刀

grad是JAX的自动微分工具。你给它一个函数，它返回这个函数的梯度函数。就这么简单，但威力巨大。

我刚开始用JAX时，最惊艳的就是grad的灵活性。它不像PyTorch那样需要你维护一个计算图，而是纯函数式的——输入函数，输出梯度函数，干净利落。

from jax import grad

def f(x):
    return x ** 3 + 2 * x ** 2 - 5 * x + 3

# 一阶导数
f_prime = grad(f)
print(f_prime(2.0))  # 3*4 + 4*2 - 5 = 15.0

# 二阶导数
f_double_prime = grad(grad(f))
print(f_double_prime(2.0))  # 6*2 + 4 = 16.0

# 高阶导数？没问题
f_triple_prime = grad(grad(grad(f)))
print(f_triple_prime(2.0))  # 6.0

grad默认对第一个参数求导。如果你需要对特定参数求导，用argnums指定：

def loss(params, inputs, targets):
    predictions = jnp.dot(inputs, params['w']) + params['b']
    return jnp.mean((predictions - targets) ** 2)

# 只对'w'求导
grad_w = grad(loss, argnums=0)
# 或者对多个参数求导
grad_w_b = grad(loss, argnums=(0, 1))

4.4 三者组合：真正的威力

单独看vmap、pmap、grad，每个都很强。但把它们组合起来，才是JAX真正的杀手锏。你想想看：自动向量化 + 多设备并行 + 自动求导，这不就是训练大模型的标配吗？

来看一个完整的训练循环示例：

from jax import grad, vmap, pmap
import jax.numpy as jnp
import jax

# 1. 定义单样本的损失
def single_loss(params, x, y):
    pred = jnp.dot(params['w'], x) + params['b']
    return (pred - y) ** 2

# 2. 用vmap向量化batch
batch_loss = vmap(single_loss, in_axes=(None, 0, 0))

# 3. 定义梯度函数
def compute_grad(params, batch_x, batch_y):
    return grad(lambda p: jnp.mean(batch_loss(p, batch_x, batch_y)))(params)

# 4. 用pmap做数据并行
parallel_grad = pmap(compute_grad, in_axes=(None, 0, 0))

# 5. 训练循环
def train_step(params, batch_x, batch_y, lr=0.01):
    grads = parallel_grad(params, batch_x, batch_y)
    # 跨设备聚合梯度
    grads = jax.lax.pmean(grads, axis_name='devices')
    # 更新参数
    params = jax.tree_map(lambda p, g: p - lr * g, params, grads)
    return params

# 初始化参数和数据
params = {'w': jnp.ones((10,)), 'b': jnp.zeros(())}
batch_x = jnp.ones((32, 10))  # 32个样本，每个10维
batch_y = jnp.ones((32,))

# 执行一步训练
params = train_step(params, batch_x, batch_y)

这里有个细节要注意：pmap内部用了axis_name='devices'，这是为了配合pmean做跨设备规约。如果你忘了加这个，每个设备只会更新自己的那部分参数，训练结果会乱掉。我曾经因为这个bug调了一整天，最后发现是少了个参数名。

4.5 性能调优的实战经验

最后分享几个我在项目中积累的经验：

• vmap vs 手动batch：vmap不是万能的。对于非常简单的操作（比如逐元素运算），手动batch可能更快。vmap的优势在于处理复杂的、有数据依赖的逻辑。
• pmap的batch size：每个设备的batch size不要太小。我一般保证每个设备至少处理16-32个样本，否则设备间的通信开销会吃掉并行收益。
• grad的编译：grad返回的函数可以配合jit使用。我习惯把grad和jit组合：jit(grad(loss))，这样梯度计算和编译优化一步到位。
• 三者嵌套：你可以把vmap放在pmap里面，也可以反过来。但要注意，pmap内部用vmap时，vmap的batch维度会被pmap切分。我建议先在小规模上验证逻辑正确性，再上大规模。

好了，vmap、pmap、grad这三个核心转换就讲完了。它们各自解决一个问题，组合起来能解决更复杂的问题。你想想看，从单样本到batch，从单设备到多设备，从函数值到梯度，JAX用这三个工具就把整个训练流程串起来了。这就是函数式编程的魅力——小而美的原语，组合出强大的系统。