具有多核 CPU 的 JAX pmap

Question

使用多CPU核心的正确方法是什么jax.pmap？

以下示例在 CPU 核心后端上为 SPMD 创建环境变量，测试 JAX 是否识别设备，并尝试设备锁定。

import os
os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=2'

import jax as jx
import jax.numpy as jnp

jx.local_device_count()
# WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
# 2

jx.devices("cpu")
# [CpuDevice(id=0), CpuDevice(id=1)]

def sfunc(x): while True: pass

jx.pmap(sfunc)(jnp.arange(2))

从jupyter内核执行并观察htop发现只有一个核心被锁定

htop当省略前两行并运行时，我收到相同的输出：

$ env XLA_FLAGS=--xla_force_host_platform_device_count=2 python test.py

替换sfunc为

def sfunc(x): return 2.0*x

并打电话

jx.pmap(sfunc)(jnp.arange(2))
# ShardedDeviceArray([0., 2.], dtype=float32, weak_type=True)

确实返回一个SharedDeviecArray.

显然我没有正确配置 JAX/XLA 以使用两个核心。我缺少什么以及我可以做什么来诊断问题？

Answer 1

据我所知，您正在正确配置核心（参见问题＃2714）。问题出在你的测试函数上：

def sfunc(x): while True: pass

该函数在跟踪时陷入无限循环，而不是在运行时。跟踪发生在单个 CPU 上的主机 Python 进程中（有关在 JAX 转换中跟踪的想法的介绍，请参阅如何在 JAX 中思考）。

如果您想在运行时观察 CPU 使用情况，则必须使用一个完成跟踪并开始运行的函数。为此，您可以使用任何实际产生结果的长时间运行的函数。这是一个简单的例子：

def sfunc(x):
  for i in range(100):
    x = (x @ x)
  return x

jx.pmap(sfunc)(jnp.zeros((2, 1000, 1000)))