JAX多GPU不同分片方式下矩阵乘法结果与性能差异问询

我正在进行JAX多GPU不同数据分片方式的矩阵乘法教程实验，发现不同分片方式不仅计算耗时不同，结果也存在细微差异。

观察结果 #

• 方法1的结果与方法0完全一致，但与方法2不同；
• 计算速度：方法1最快，其次是方法2，方法0最慢。

方法说明 #

方法0 #

仅在单GPU设备上执行矩阵乘法。

方法1 #

x = jax.random.normal(jax.random.PRNGKey(0), (8192, 8192))
y = jax.device_put(x, sharding.reshape(4, 2).replicate(1))
z = jax.device_put(x, sharding.reshape(4, 2).replicate(0))
print('lhs sharding:')
jax.debug.visualize_array_sharding(y)
print('rhs sharding:')
jax.debug.visualize_array_sharding(z)

w = jnp.dot(y, z)

方法2 #

x = jax.random.normal(jax.random.PRNGKey(0), (8192, 8192))
y = jax.device_put(x, sharding.reshape(4, 2))
z = jax.device_put(x, sharding.reshape(4, 2))
print('lhs sharding:')
jax.debug.visualize_array_sharding(y)
print('rhs sharding:')
jax.debug.visualize_array_sharding(z)

w = jnp.dot(y, z)

结果解释与疑问解答 #

1. 观察结果的原因 #

结果一致性差异

方法0是单GPU计算，所有浮点运算的执行顺序固定，结果唯一。方法1中，y在列维度（第1维）全复制，z在行维度（第0维）全复制：每个GPU都持有y的完整列分片和z的完整行分片，子矩阵乘法的计算逻辑和单GPU完全一致，最终聚合结果时不会因为浮点运算顺序变化产生误差，因此和方法0结果完全相同。

方法2中，y和z都在二维设备网格上分片，矩阵乘法过程需要跨设备进行All-Gather、Reduce-Scatter等通信操作，不同GPU上的子块计算顺序与单GPU不同。由于浮点运算的固有特性，不同的累积顺序会产生微小的舍入误差，导致结果与方法0、方法1出现细微差异。

计算速度差异

方法1的分片方式完美匹配矩阵乘法的并行逻辑：每个GPU可以独立完成对应子矩阵的乘法计算，几乎不需要跨设备通信，充分利用多GPU的并行算力，因此速度最快。

方法2采用二维分片，矩阵乘法过程中需要频繁跨设备交换中间数据，通信开销会抵消部分并行计算的收益，所以速度比方法1慢，但仍远快于单GPU的方法0。

方法0无并行加速，完全依赖单GPU算力，因此速度最慢。

2. 主流机器学习算法的一致性处理 #

主流深度学习框架（如JAX、PyTorch、TensorFlow）在分布式训练中会自动处理一致性问题：

• 数据并行场景下，框架会自动执行梯度的All-Reduce操作，保证所有GPU的参数更新完全一致；
• 模型并行场景下，框架会封装跨设备的通信与计算逻辑，尽量维持数值稳定性，但浮点运算的细微舍入误差几乎无法完全避免——不过深度学习模型对这类微小误差鲁棒性极强，通常不会影响最终训练效果；
• 若对确定性有严格要求，可通过启用高精度浮点（如JAX的jax.enable_x64()）或指定运算顺序来减少差异，但可能会牺牲部分计算性能。

内容的提问来源于Stack Exchange，提问作者Dmitry J