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, compiler_options=None)[source]#
使
fun
被编译并自动在多个设备上分区。注意:此函数现在等同于 jax.jit,请改用该函数。返回的函数具有与
fun
语义等效的语义,但被编译为跨多个设备(例如,多个 GPU 或多个 TPU 核心)运行的 XLA 计算。如果fun
的 jitted 版本不适合单个设备的内存,或者通过在多个设备上并行运行每个操作来加速fun
,则此功能可能很有用。设备上的分区自动发生,基于在
in_shardings
中指定的输入分区和在out_shardings
中指定的输出分区的传播。这两个参数中指定的资源必须引用网格轴,如jax.sharding.Mesh()
上下文管理器所定义。请注意,pjit()
应用时的网格定义将被忽略,并且返回的函数将使用每个调用站点可用的网格定义。如果
pjit()
'd 函数的输入尚未基于in_shardings
正确分区,则它们将自动跨设备分区。在某些情况下,确保输入已正确预分区可以提高性能。例如,如果将一个pjit()
'd 函数的输出传递给另一个pjit()
'd 函数(或循环中的同一个pjit()
'd 函数),请确保相关的out_shardings
与相应的in_shardings
匹配。注意
多进程平台: 在多进程平台(如 TPU pod)上,
pjit()
可用于跨进程在所有可用设备上运行计算。为了实现这一点,pjit()
被设计用于 SPMD Python 程序中,其中每个进程都运行相同的 Python 代码,以便所有进程都以相同的顺序运行相同的pjit()
'd 函数。在此配置中运行时,网格应包含跨所有进程的设备。所有输入参数都必须是全局形状。
fun
仍将在网格中的所有设备(包括来自其他进程的设备)上执行,并将获得跨多个进程分布的数据的全局视图,作为一个单独的数组。SPMD 模型还要求相同的多进程
pjit()
'd 函数必须在所有进程上以相同的顺序运行,但它们可以与在单个进程中运行的任意操作交错。- 参数:
fun (Callable) – 要编译的函数。应为纯函数,因为副作用可能只执行一次。其参数和返回值应为数组、标量或(嵌套)标准 Python 容器(tuple/list/dict)。由
static_argnums
指示的位置参数可以是任何内容,只要它们是可哈希的并且定义了相等操作。静态参数包含在编译缓存键的一部分中,这就是为什么必须定义哈希和相等运算符。in_shardings (Any) –
结构与
fun
的参数结构匹配的 Pytree,所有实际参数都替换为资源分配规范。也可以指定 Pytree 前缀(例如,一个值代替整个子树),在这种情况下,叶子将广播到该子树中的所有值。in_shardings
参数是可选的。JAX 将从输入jax.Array
's 推断分片,如果无法推断分片,则默认为复制输入。有效的资源分配规范为
Sharding
,它将决定如何对值进行分区。使用它,不需要网格上下文管理器。None
是一种特殊情况,其语义为如果未提供网格上下文管理器,则 JAX 可以自由选择它想要的任何分片。对于 in_shardings,JAX 将其标记为复制,但此行为将来可能会更改。对于 out_shardings,我们将依赖 XLA GSPMD 分区器来确定输出分片。
如果提供了网格上下文管理器,则 None 将意味着该值将在网格的所有设备上复制。
为了向后兼容,in_shardings 仍然支持摄取
PartitionSpec
。此选项只能与网格上下文管理器一起使用。PartitionSpec
,一个元组,其长度最多等于分区值的秩。每个元素可以是None
、网格轴或网格轴元组,并指定分配给分区值维度(与其在规范中的位置匹配)的资源集。
每个维度的大小都必须是分配给它的资源总数的倍数。
out_shardings (Any) – 类似于
in_shardings
,但为函数输出指定资源分配。out_shardings
参数是可选的。如果未指定,jax.jit()
将使用 GSPMD 的分片传播来确定如何对输出进行分片。static_argnums (int | Sequence[int] | None | None) –
一个可选的整数或整数集合,用于指定哪些位置参数应被视为静态(编译时常量)。仅依赖于静态参数的操作将在 Python 中(在跟踪期间)进行常量折叠,因此相应的参数值可以是任何 Python 对象。
静态参数应该是可哈希的,这意味着实现了
__hash__
和__eq__
,并且是不可变的。使用这些常量的不同值调用 jitted 函数将触发重新编译。不是数组或其容器的参数必须标记为静态。如果未提供
static_argnums
,则没有参数被视为静态。static_argnames (str | Iterable[str] | None | None) – 一个可选的字符串或字符串集合,用于指定哪些命名参数应被视为静态(编译时常量)。有关详细信息,请参阅关于
static_argnums
的注释。如果未提供但设置了static_argnums
,则默认值基于调用inspect.signature(fun)
以查找相应的命名参数。donate_argnums (int | Sequence[int] | None | 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 | None) – 一个可选的字符串或字符串集合,用于指定哪些命名参数被捐赠给计算。有关详细信息,请参阅关于
donate_argnums
的注释。如果未提供但设置了donate_argnums
,则默认值基于调用inspect.signature(fun)
以查找相应的命名参数。keep_unused (bool) – 如果 False(默认值),则 JAX 确定 fun 未使用的参数可能会从生成的已编译 XLA 可执行文件中删除。此类参数将不会传输到设备,也不会提供给底层可执行文件。如果为 True,则不会修剪未使用的参数。
device (xc.Device | None | None) – 此参数已弃用。请在将参数传递给 jit 之前将其放在您想要的设备上。可选,jitted 函数将在其上运行的设备。(可用设备可以通过
jax.devices()
检索。)默认值继承自 XLA 的 DeviceAssignment 逻辑,通常是使用jax.devices()[0]
。backend (str | None | None) – 此参数已弃用。请在将参数传递给 jit 之前将其放在您想要的后端上。可选,表示 XLA 后端的字符串:
'cpu'
、'gpu'
或'tpu'
。inline (bool)
abstracted_axes (Any | None | 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. ]