jax.pmap#

jax.pmap(fun, axis_name=None, *, in_axes=0, out_axes=0, static_broadcasted_argnums=(), devices=None, backend=None, axis_size=None, donate_argnums=(), global_arg_shapes=None)[源代码]#

支持集体操作的并行映射。

注意

pmap() 现在基于 jit() 和 shard_map() 实现。请参阅 迁移指南 以获取更多信息。

pmap() 的目的是表达单程序多数据 (SPMD) 程序。将 pmap() 应用于函数将使用 XLA 进行编译（类似于 jit()），然后将其并行执行在 XLA 设备上，例如多个 GPU 或多个 TPU 核心。语义上它可与 vmap() 相比，因为两种变换都会在数组轴上映射函数，但 vmap() 通过将映射的轴推送到原始操作来向量化函数，而 pmap() 而是复制函数并在每个副本上并行执行在自己的 XLA 设备上。

映射轴的大小必须小于或等于可用的本地 XLA 设备数量，由 jax.local_device_count() 返回（除非指定了 devices，见下文）。对于嵌套的 pmap() 调用，映射轴大小的乘积必须小于或等于 XLA 设备数量。

注意

pmap() 会编译 fun，因此虽然它可以与 jit() 结合使用，但通常是不必要的。

pmap() 要求所有参与的设备都相同。例如，无法使用 pmap() 将计算并行化到两个不同型号的 GPU 上。目前，同一个设备在同一个 pmap 中参与两次是错误的。

多进程平台：在多进程平台（如 TPU Pods）上，pmap() 被设计用于 SPMD Python 程序，其中每个进程运行相同的 Python 代码，从而所有进程以相同的顺序运行相同的 pmapped 函数。每个进程仍应使用映射轴大小等于*本地*设备数量（除非指定了 devices，见下文）来调用 pmapped 函数，并且将像往常一样返回具有相同领先轴大小的数组。但是，fun 中的任何集体操作都将在*所有*参与的设备（包括其他进程的设备）上进行计算，通过设备间通信。概念上，这可以被视为在一个跨进程分片的单个数组上运行 pmap，其中每个进程“只看到”输入和输出的本地分片。SPMD 模型要求相同的多进程 pmap 必须在所有设备上以相同的顺序运行，但它们可以与在单个进程中运行的任意操作交错。

参数:

  • fun (Callable) – 要映射到参数轴上的函数。其参数和返回值应该是数组、标量，或者它们的（嵌套）标准 Python 容器（元组/列表/字典）。由 static_broadcasted_argnums 指定的位置参数可以是任何东西，只要它们是可哈希的并且定义了相等操作。

  • axis_name (AxisName | None) – 可选，一个用于标识映射轴的可哈希 Python 对象，以便可以应用并行集体操作。

  • in_axes (int | None | Sequence[Any]) – 一个非负整数、None 或它们的嵌套 Python 容器，用于指定要映射哪些位置参数的轴。作为关键字参数传递的参数总是映射到它们的领先轴（即轴索引 0）。有关详细信息，请参阅 vmap()。

  • out_axes (Any) – 一个非负整数、None 或它们的嵌套 Python 容器，指示映射轴应出现在输出的哪个位置。所有具有映射轴的输出都必须具有非 None 的 out_axes 规范（参见 vmap()）。

  • static_broadcasted_argnums (int | Iterable[int]) –

    一个整数或一组整数，用于指定哪些位置参数应被视为静态（编译时常量）。仅依赖于静态参数的操作将进行常量折叠。使用这些常量的不同值调用 pmapped 函数将触发重新编译。如果使用少于 static_broadcasted_argnums 指定的位置参数调用 pmapped 函数，则会引发错误。每个静态参数都将广播到所有设备。不是数组或其容器的参数必须标记为静态。默认为 ()。

    静态参数必须是可哈希的，这意味着 __hash__ 和 __eq__ 都已实现，并且应该是不可变的。

  • devices (Sequence[xc.Device] | None) – 这是一个实验性功能，API 可能会发生变化。可选，一个要映射的设备序列。（可通过 jax.devices() 检索可用设备）。在多进程设置中，每个进程必须提供相同的参数（因此将包括跨进程的设备）。如果指定，映射轴的大小必须等于给定进程本地设备序列的数量。嵌套的 pmap() 在内层或外层 pmap() 中指定 devices 尚不支持。

  • backend (str | None) – 这是一个实验性功能，API 可能会发生变化。可选，一个表示 XLA 后端的字符串。‘cpu’、‘gpu’ 或 ‘tpu’。

  • axis_size (int | None) – 可选；映射轴的大小。

  • donate_argnums (int | Iterable[int]) –

    指定哪些位置参数缓冲区被“捐赠”给计算。如果您不再需要参数缓冲区，在计算完成后捐赠它们是安全的。在某些情况下，XLA 可以利用捐赠的缓冲区来减少执行计算所需的内存量，例如，回收您的一块输入缓冲区来存储结果。您不应重用您捐赠给计算的缓冲区，JAX 会在您尝试时引发错误。请注意，donate_argnums 仅适用于位置参数，关键字参数不会被捐赠。

    有关缓冲区捐赠的更多详细信息，请参阅 FAQ。

  • global_arg_shapes (tuple[tuple[int, ...], ...] | None)

返回:

一个 fun 的并行化版本，其参数与 fun 对应，但在 in_axes 指定的位置具有额外的数组轴，并且输出具有一个额外的领先数组轴（大小相同）。

返回类型:

任意类型

例如，假设有 8 个 XLA 设备可用，pmap() 可以用作沿着领先数组轴的映射

>>> import jax.numpy as jnp
>>>
>>> out = pmap(lambda x: x ** 2)(jnp.arange(8))  
>>> print(out)  
[0, 1, 4, 9, 16, 25, 36, 49]

当领先维度小于可用设备数量时，JAX 将只在部分设备上运行

>>> x = jnp.arange(3 * 2 * 2.).reshape((3, 2, 2))
>>> y = jnp.arange(3 * 2 * 2.).reshape((3, 2, 2)) ** 2
>>> out = pmap(jnp.dot)(x, y)  
>>> print(out)  
[[[    4.     9.]
  [   12.    29.]]
 [[  244.   345.]
  [  348.   493.]]
 [[ 1412.  1737.]
  [ 1740.  2141.]]]

如果您的领先维度大于可用设备数量，您将收到一个错误

>>> pmap(lambda x: x ** 2)(jnp.arange(9))  
ValueError: ... requires 9 replicas, but only 8 XLA devices are available

与 vmap() 一样，在 in_axes 中使用 None 表示一个参数没有额外的轴，应该被广播而不是映射到副本上

>>> x, y = jnp.arange(2.), 4.
>>> out = pmap(lambda x, y: (x + y, y * 2.), in_axes=(0, None))(x, y)  
>>> print(out)  
([4., 5.], [8., 8.])

请注意，pmap() 总是返回在领先轴上映射的值，相当于在 vmap() 中使用 out_axes=0。

除了表达纯粹的映射之外，pmap() 还可以用来表达通过集体操作进行通信的并行单程序多数据 (SPMD) 程序。例如

>>> f = lambda x: x / jax.lax.psum(x, axis_name='i')
>>> out = pmap(f, axis_name='i')(jnp.arange(4.))  
>>> print(out)  
[ 0.          0.16666667  0.33333334  0.5       ]
>>> print(out.sum())  
1.0

在此示例中，axis_name 是一个字符串，但它可以是任何具有 __hash__ 和 __eq__ 定义的 Python 对象。

pmap() 的参数 axis_name 为映射轴命名，以便集体操作（如 jax.lax.psum()）可以引用它。轴名称在嵌套 pmap() 函数的情况下尤为重要，集体操作可以作用于不同的轴

>>> from functools import partial
>>> import jax
>>>
>>> @partial(pmap, axis_name='rows')
... @partial(pmap, axis_name='cols')
... def normalize(x):
...   row_normed = x / jax.lax.psum(x, 'rows')
...   col_normed = x / jax.lax.psum(x, 'cols')
...   doubly_normed = x / jax.lax.psum(x, ('rows', 'cols'))
...   return row_normed, col_normed, doubly_normed
>>>
>>> x = jnp.arange(8.).reshape((4, 2))
>>> row_normed, col_normed, doubly_normed = normalize(x)  
>>> print(row_normed.sum(0))  
[ 1.  1.]
>>> print(col_normed.sum(1))  
[ 1.  1.  1.  1.]
>>> print(doubly_normed.sum((0, 1)))  
1.0

在多进程平台上，集体操作将在所有设备上运行，包括其他进程的设备。例如，假设以下代码在具有 4 个 XLA 设备（每个）的两个进程上运行

>>> f = lambda x: x + jax.lax.psum(x, axis_name='i')
>>> data = jnp.arange(4) if jax.process_index() == 0 else jnp.arange(4, 8)
>>> out = pmap(f, axis_name='i')(data)  
>>> print(out)  
[28 29 30 31] # on process 0
[32 33 34 35] # on process 1

每个进程传入一个长度为 4 的不同数组，对应于其 4 个本地设备，psum 操作将作用于所有 8 个值。概念上，这两个长度为 4 的数组可以被视为一个分片的长度为 8 的数组（在此示例中等同于 jnp.arange(8)），对其进行映射，映射轴的长度为 8，命名为 ‘i’。每个进程上的 pmap 调用然后返回相应的长度为 4 的输出分片。

可以通过 devices 参数来指定确切用于运行并行计算的设备。例如，再次假设一个具有 8 个设备的单个进程，以下代码定义了两个并行计算，一个在最初的六个设备上运行，另一个在剩余的两个设备上运行

>>> from functools import partial
>>> @partial(pmap, axis_name='i', devices=jax.devices()[:6])
... def f1(x):
...   return x / jax.lax.psum(x, axis_name='i')
>>>
>>> @partial(pmap, axis_name='i', devices=jax.devices()[-2:])
... def f2(x):
...   return jax.lax.psum(x ** 2, axis_name='i')
>>>
>>> print(f1(jnp.arange(6.)))  
[0.         0.06666667 0.13333333 0.2        0.26666667 0.33333333]
>>> print(f2(jnp.array([2., 3.])))  
[ 13.  13.]