JAX vmap、pmap与Python multiprocessing对比及迁移疑问

原Python并行实现代码 #

# start pool process 
pool = multiprocessing.Pool(processes=10) # if node has 10 CPU cores, start 10 processes

# use pool.map to evaluate function(input) for each input in parallel
# suppose len(inputs) is very large and 10 inputs are processed in parallel at a time
# store the results in a list called out
out = pool.map(function,inputs)

# close pool processes to free memory
pool.close()
pool.join()

疑问与解答 #

1. vmap(function,in_axes=0)(inputs)是否会分配到所有可用CPU核心？

   • 会。vmap通过向量化转换目标函数，交由XLA自动调度，以单进程多线程的方式利用所有CPU核心，进程内并行的开销远低于multiprocessing的进程级并行。

2. pmap(function,in_axes=0)(inputs)与vmap、multiprocessing.pool.map有何区别？

   • vmap：单进程内的向量化并行，由XLA做底层优化，适合轻量循环场景，无额外进程开销，仅利用单进程的多线程资源（覆盖所有CPU核心）。
   • pmap：跨XLA设备的SPMD并行，每个设备对应独立进程，要求输入按设备数量拆分。并行粒度为进程级，但绑定XLA识别的逻辑设备，而非手动创建的进程池。
   • multiprocessing.pool.map：手动创建进程池实现进程级并行，输入拆分由Python调度，不依赖XLA，对非JAX兼容的Python代码兼容性更强，但进程创建和调度的开销高于vmap。

3. 上述multiprocessing.pool.map的用法是否属于pmap适用的“单程序多数据（SPMD）”代码示例？

   • 属于SPMD范畴：所有执行单元（进程）运行相同的function，处理不同的inputs分片。但pmap的SPMD要求输入必须按XLA设备数均匀拆分，而multiprocessing的SPMD对输入拆分的灵活性更高，无需严格匹配核心/设备数量。

4. 执行pmap(function,in_axes=0)(inputs)时出现错误：ValueError: compiling computation that requires 10 logical devices, but only 1 XLA devices are available (num_replicas=10, num_partitions=1)，这是什么意思？

   • 该错误表示pmap试图创建10个计算副本（通常对应输入的批量维度或手动指定的num_replicas参数），但当前JAX仅识别到1个CPU逻辑设备。默认情况下，JAX将整个CPU视为单个逻辑设备，不会自动拆分核心。要解决这个问题，需手动配置JAX的CPU逻辑设备数，例如：

     import jax
     jax.config.update("jax_cpu_core_count", 10)
     

5. 仅需利用单节点的部分/全部CPU核心（如MacBook的10个核心）。嵌套pmap(vmap)是否用于多节点（如超级计算机）核心的并行？这更类似mpi4py而非multiprocessing？

   • 嵌套pmap(vmap)主要用于多节点+单节点多设备的层级并行场景：外层pmap负责跨节点的通信与并行（类似mpi4py的多节点调度），内层vmap则在单节点内的设备上做向量化并行。如果仅需单节点利用全部核心，直接使用vmap即可（XLA会自动调度到所有核心）；若要使用pmap，只需手动配置单节点的逻辑CPU设备数，无需涉及多节点。multiprocessing仅限单节点，而mpi4py支持跨节点，嵌套pmap的外层逻辑确实更接近mpi4py的应用场景。

内容的提问来源于stack exchange，提问作者Jim Raynor