JAX多GPU并行训练PINN时利用率低、速度慢问题咨询

我尝试用4块GPU求解物理信息神经网络（PINN）问题。单GPU训练时，GPU利用率能到100%，训练速度200 [it/s]；但用JAX自动并行化的分片+jax.jit策略（参考官方文档），把数据分片到不同GPU、复制模型参数与状态后，4块GPU利用率不足10%，训练速度仅20 [it/s]。我也试过jax.pmap方法（参考相关教程），代码示例如下：

import functools

# Remember that the 'G' is just an arbitrary string label used
#          to later tell 'jax.lax.pmean' which axis to reduce over.
# Here, we call it
#          'G', but could have used anything,
#          so long as 'pmean' used the same.

@functools.partial(jax.pmap, axis_name="G")
def update(params: Params, x: jnp.ndarray, y: jnp.ndarray):
    # Compute the gradients on the given minibatch (individually on each device)
    loss, grads = jax.value_and_grad(loss_fn)(params, x, y)

    # Combine the gradient across all devices (by taking their mean)
    grads = jax.lax.pmean(grads, axis_name="G")

    # Also combine the loss. Unnecessary for the update, but useful for logging
    loss = jax.lax.pmean(loss, axis_name="G")

    # Each device performs its own update, but since we start with the same params
    # and synchronise gradients, the params stay in sync
    LEARNING_RATE = 1e-3
    new_params = jax.tree_map(
       lambda param, g: param - g * LEARNING_RATE, params, grads)
    return new_params, loss

以下是可能的原因：

• 单设备batch过小，通信开销占比过高：PINN的loss计算涉及大量微分运算（如自动微分求解PDE残差），若分给每个GPU的batch size未随GPU数量等比例放大，单设备计算量会远小于跨GPU梯度同步（pmean操作）的通信开销，导致GPU大部分时间处于等待通信的空闲状态，利用率暴跌。
• PINN计算图未被JAX并行策略正确适配：PINN的loss函数通常包含复杂的自定义微分逻辑（如对输入变量多次求导），jax.jit或pmap可能未正确捕获这些计算，生成的计算图效率低下；若loss函数中存在未被JAX追踪的Python控制流，还会触发频繁JIT编译或计算回退到CPU，大幅拉低GPU利用率。
• 数据分片与参数复制存在额外开销：若数据在CPU上分片后再传输到GPU，而非直接在设备端完成分片，会产生不必要的主机-设备数据拷贝；参数复制时若出现冗余的数据传输，也会导致GPU等待数据加载，无法满负荷运行。
• GPU通信带宽瓶颈：若GPU之间仅通过PCIe连接（无NVLink），跨设备通信带宽远低于设备内部计算带宽。PINN的梯度包含大量参数的梯度值，同步时的数据量若超出通信带宽承载能力，会导致GPU长时间等待，利用率上不去。
• 训练循环中的同步操作拖慢计算：若每次迭代都执行主机端日志打印、loss从GPU拷贝到CPU等同步操作，会打断GPU的连续计算流程，导致GPU频繁空闲。比如代码中返回的loss若每次都同步拷贝到CPU打印，未做异步处理，会显著降低整体训练速度。

内容的提问来源于stack exchange，提问作者WANG Jacques