使用多核CPU的JAX PMAP

Question

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

以下示例为CPU Core后端上的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 显示，只有一个核心锁定

当省略前两行和运行：

$ 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

据我所知，您正在正确配置核心（请参阅eg 问题＃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)))