jax.experimental.pjit 模块

目录

jax.experimental.pjit 模块#

API#

jax.experimental.pjit.pjit(fun, in_shardings=UnspecifiedValue, out_shardings=UnspecifiedValue, static_argnums=None, static_argnames=None, donate_argnums=None, donate_argnames=None, keep_unused=False, device=None, backend=None, inline=False, abstracted_axes=None)[源代码][源代码]#

使 fun 被编译并在多个设备上自动分区。

注意：此函数现在等同于 jax.jit，请改用该函数。返回的函数在语义上等同于 fun，但编译为跨多个设备（例如多个 GPU 或多个 TPU 核心）运行的 XLA 计算。如果 fun 的即时编译版本不适合单个设备的内存，或者通过在多个设备上并行运行每个操作来加速 fun，这可能很有用。

设备上的分区是根据 in_shardings 中指定的输入分区和 out_shardings 中指定的输出分区自动进行的。这两个参数中指定的资源必须引用由 jax.sharding.Mesh() 上下文管理器定义的网格轴。请注意，在 pjit() 应用时的网格定义将被忽略，返回的函数将使用每个调用站点可用的网格定义。

如果输入尚未根据 in_shardings 正确分区，则传递给 pjit() 函数的输入将自动跨设备分区。在某些情况下，确保输入已经正确预分区可以提高性能。例如，如果将一个 pjit() 函数的输出传递给另一个 pjit() 函数（或在循环中传递给同一个 pjit() 函数），请确保相关的 out_shardings 与相应的 in_shardings 匹配。

备注

多进程平台: 在多进程平台如 TPU pods 上，pjit() 可用于跨所有可用设备跨进程运行计算。为此，pjit() 设计用于 SPMD Python 程序中，其中每个进程都运行相同的 Python 代码，以便所有进程按相同顺序运行相同的 pjit() 函数。

在这种配置下运行时，网格应包含所有进程中的设备。所有输入参数必须是全局形状的。fun 仍将在网格中的所有设备上执行，包括来自其他进程的设备，并将获得跨多个进程分布的数据的全局视图，作为一个单一数组。

SPMD 模型还要求在所有进程上必须以相同的顺序运行相同的 pjit() 函数，但它们可以与在单个进程中运行的任意操作交错。

参数:

  • fun (Callable) – 要编译的函数。应为纯函数，因为副作用可能只会执行一次。其参数和返回值应为数组、标量或（嵌套的）标准 Python 容器（元组/列表/字典）。由 static_argnums 指示的位置参数可以是任何内容，前提是它们是可哈希的并且定义了相等操作。静态参数作为编译缓存键的一部分包含在内，这就是为什么必须定义哈希和相等运算符的原因。

  • in_shardings –

    与 fun 参数结构匹配的 Pytree，其中所有实际参数都被资源分配规范替换。也可以指定一个 pytree 前缀（例如，用一个值代替整个子树），在这种情况下，叶子会被广播到该子树中的所有值。

    in_shardings 参数是可选的。JAX 会从输入的 jax.Array 推断分片，如果无法推断分片，则默认复制输入。

    有效的资源分配规范是：

    • Sharding，它将决定值如何被分区。有了这个，使用网格上下文管理器就不再是必需的。

    • None is a special case whose semantics are:

      • 如果未提供网格上下文管理器，JAX可以自由选择任何它想要的拆分方式。对于in_shardings，JAX会将其标记为复制，但此行为未来可能会改变。对于out_shardings，我们将依赖XLA GSPMD分区器来确定输出拆分。

      • 如果提供了网格上下文管理器，None 将意味着该值将在网格的所有设备上复制。

    • 为了向后兼容，in_shardings 仍然支持摄取 PartitionSpec。此选项只能与网格上下文管理器一起使用。

      • PartitionSpec, a tuple of length at most equal to the rank of the partitioned value. Each element can be a None, a mesh axis or a tuple of mesh axes, and specifies the set of resources assigned to partition the value’s dimension matching its position in the spec.

    每个维度的尺寸必须是分配给它的资源总数的倍数。

  • out_shardings – 类似于 in_shardings，但指定函数输出的资源分配。out_shardings 参数是可选的。如果未指定，jax.jit() 将使用 GSPMD 的分片传播来确定如何分片输出。

  • static_argnums (int | Sequence[int] | None) – 一个可选的整数或整数集合，用于指定哪些位置参数作为静态参数（编译时常量）处理。仅依赖于静态参数的操作将在Python中进行常量折叠（在跟踪期间），因此相应参数的值可以是任何Python对象。静态参数应是可哈希的，即实现了 __hash__ 和 __eq__，并且是不可变的。使用这些常量的不同值调用jitted函数将触发重新编译。不是数组或其容器的参数必须标记为静态。如果未提供 static_argnums，则不将任何参数视为静态。

  • static_argnames (str | Iterable[str] | None) – 一个可选的字符串或字符串集合，指定哪些命名参数作为静态（编译时常量）处理。详情请参阅 static_argnums 的注释。如果没有提供但 static_argnums 已设置，默认值基于调用 inspect.signature(fun) 来查找相应的命名参数。

  • donate_argnums (int | Sequence[int] | None) – 指定哪些位置参数缓冲区被“捐赠”给计算。如果你在计算完成后不再需要这些缓冲区，那么捐赠它们是安全的。在某些情况下，XLA 可以利用捐赠的缓冲区来减少执行计算所需的内存量，例如将你的一个输入缓冲区回收用于存储结果。你不应该重用捐赠给计算的缓冲区，如果你尝试这样做，JAX 会抛出错误。默认情况下，不捐赠任何参数缓冲区。如果没有提供 donate_argnums 或 donate_argnames，则不捐赠任何参数。如果未提供 donate_argnums 但提供了 donate_argnames，或者反之，JAX 使用 inspect.signature(fun) 来查找与 donate_argnames``（或反之）对应的任何位置参数。如果同时提供了 ``donate_argnums 和 donate_argnames，则不使用 inspect.signature，并且只有列在 donate_argnums 或 donate_argnames 中的实际参数会被捐赠。有关缓冲区捐赠的更多详情，请参阅 FAQ。

  • donate_argnames (str | Iterable[str] | None) – 一个可选的字符串或字符串集合，指定哪些命名参数被捐赠给计算。详情请参阅 donate_argnums 的注释。如果没有提供但 donate_argnums 已设置，默认值基于调用 inspect.signature(fun) 来查找相应的命名参数。

  • keep_unused (bool) – 如果 False`（默认值），JAX 确定为 `fun 未使用的参数 可能 会从生成的编译 XLA 可执行文件中删除。这些参数不会被传输到设备，也不会提供给底层可执行文件。如果 True，未使用的参数将不会被修剪。

  • device (xc.Device | None) – 此参数已弃用。请在传递给 jit 之前将您的参数放在您想要的设备上。可选，jit 函数将在其上运行的设备。（可用设备可以通过 jax.devices() 获取。）默认值继承自 XLA 的 DeviceAssignment 逻辑，通常是使用 jax.devices()[0]。

  • backend (str | None) – 此参数已弃用。请在将参数传递给 jit 之前，将其放在您想要的 XLA 后端上。可选，表示 XLA 后端的字符串：'cpu'、'gpu' 或 'tpu'。

  • inline (bool)

  • abstracted_axes (Any | None)

返回:

fun 的包装版本，设置为即时编译，并根据每次调用时可用的网格自动分区。

返回类型:

JitWrapped

例如，通过单个 pjit() 应用，卷积运算符可以自动在任意一组设备上进行分区：

>>> import jax
>>> import jax.numpy as jnp
>>> import numpy as np
>>> from jax.sharding import Mesh, PartitionSpec
>>> from jax.experimental.pjit import pjit
>>>
>>> x = jnp.arange(8, dtype=jnp.float32)
>>> f = pjit(lambda x: jax.numpy.convolve(x, jnp.asarray([0.5, 1.0, 0.5]), 'same'),
...         in_shardings=None, out_shardings=PartitionSpec('devices'))
>>> with Mesh(np.array(jax.devices()), ('devices',)):
...   print(f(x))  
[ 0.5  2.   4.   6.   8.  10.  12.  10. ]
目录