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)[source]#
使
fun
编译并在多个设备之间自动分区。注意:此函数现在等同于 jax.jit,请改用 jax.jit。返回的函数语义等同于
fun
,但会被编译成一个跨多个设备(例如多个GPU或多个TPU核心)运行的XLA计算。如果fun
的 jit 版本无法放入单个设备的内存中,或者希望通过在多个设备上并行运行每个操作来加速fun
,这将非常有用。设备上的分区是根据
in_shardings
中指定的输入分区和out_shardings
中指定的输出分区的传播自动发生的。这两个参数中指定的资源必须引用网格轴,如jax.sharding.Mesh()
上下文管理器定义的那样。请注意,在pjit()
应用时的网格定义会被忽略,返回的函数将在每个调用站点使用可用的网格定义。如果
pjit()
函数的输入尚未根据in_shardings
正确分区,则这些输入会自动跨设备分区。在某些情况下,确保输入已正确预分区可以提高性能。例如,如果将一个pjit()
函数的输出传递给另一个pjit()
函数(或在循环中使用相同的pjit()
函数),请确保相关的out_shardings
与相应的in_shardings
匹配。注意
多进程平台:在 TPU Pod 等多进程平台上,
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
是一个特殊情况,其语义是如果没有提供网格上下文管理器,则 JAX 可以自由选择任何分片。对于 in_shardings,JAX 将将其标记为复制,但此行为将来可能会更改。对于 out_shardings,我们将依靠 XLA GSPMD 分区器来确定输出分片。
如果提供了网格上下文管理器,则 None 表示该值将在网格的所有设备上复制。
为了向后兼容,in_shardings 仍然支持摄取
PartitionSpec
。此选项只能与网格上下文管理器一起使用。PartitionSpec
,一个长度最多等于被分区值的秩的元组。每个元素可以是None
、网格轴或网格轴的元组,并指定分配给分区与规范中位置匹配的值维度的资源集。
每个维度的尺寸必须是分配给它的资源总数的倍数。
out_shardings – 与
in_shardings
类似,但指定函数输出的资源分配。out_shardings
参数是可选的。如果未指定,则jax.jit()
将使用 GSPMD 的分片传播来确定如何对输出进行分片。static_argnums (int | Sequence[int] | None | None) –
一个可选的整数或整数集合,指定哪些位置参数应视为静态(编译时常量)。仅依赖于静态参数的操作将在 Python 中进行常量折叠(在跟踪期间),因此相应的参数值可以是任何 Python 对象。
静态参数应该是可散列的,这意味着实现了
__hash__
和__eq__
,并且是不可变的。使用这些常量的不同值调用 jit 函数将触发重新编译。不是数组或其容器的参数必须标记为静态。如果未提供
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 之前将其放置在您想要的设备上。可选,jit 函数将在其上运行的设备。(可用的设备可以通过
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. ]