jax.Array 迁移

yashkatariya@

TL;DR

从 0.4.1 版本开始，JAX 将其默认数组实现切换为新的 jax.Array。本指南解释了此举背后的原因、可能对您的代码产生的影响以及如何（临时）切换回旧的行为。

发生了什么？

jax.Array 是一种统一的数组类型，它包含了 JAX 中的 DeviceArray、ShardedDeviceArray 和 GlobalDeviceArray 类型。jax.Array 类型有助于将并行性作为 JAX 的核心特性，简化并统一 JAX 内部结构，并允许我们统一 jit 和 pjit。如果您的代码没有提及 DeviceArray 与 ShardedDeviceArray 与 GlobalDeviceArray，则无需进行任何更改。但是，依赖于这些单独类的细节的代码可能需要进行调整才能使用统一的 jax.Array。

迁移完成后，jax.Array 将成为 JAX 中唯一的数组类型。

本文档解释了如何将现有代码库迁移到 jax.Array。有关使用 jax.Array 和 JAX 并行 API 的更多信息，请参阅分布式数组和自动并行化教程。

如何启用 jax.Array？

您可以通过以下方式启用 jax.Array：

• 将 shell 环境变量 JAX_ARRAY 设置为类似 true 的值（例如，1）；

• 如果您的代码使用 absl 解析标志，则将布尔标志 jax_array 设置为类似 true 的值；

• 在您的主文件的顶部使用以下语句

  import jax
  jax.config.update('jax_array', True)
  

如何知道 jax.Array 是否破坏了我的代码？

判断 jax.Array 是否导致任何问题的最简单方法是禁用 jax.Array 并查看问题是否消失。

现在如何禁用 jax.Array？

在 **2023 年 3 月 15 日** 之前，可以通过以下方式禁用 jax.Array：

• 将 shell 环境变量 JAX_ARRAY 设置为类似 false 的值（例如，0）；

• 如果您的代码使用 absl 解析标志，则将布尔标志 jax_array 设置为类似 false 的值；

• 在您的主文件的顶部使用以下语句

  import jax
  jax.config.update('jax_array', False)
  

为什么要创建 jax.Array？

目前 JAX 有三种类型：DeviceArray、ShardedDeviceArray 和 GlobalDeviceArray。jax.Array 合并了这三种类型，并在添加新的并行功能的同时清理了 JAX 的内部结构。

我们还引入了一个新的 Sharding 抽象，它描述了一个逻辑 Array 如何在多个设备（例如 TPU 或 GPU）上进行物理分片。此更改还升级、简化并将 pjit 的并行功能合并到 jit 中。用 jit 装饰的函数将能够在分片数组上运行，而无需将数据复制到单个设备。

使用 jax.Array 可以获得的功能

• C++ pjit 分派路径

• 逐操作并行性（即使数组分布在多个主机上的多个设备上）

• 使用 pjit/jit 进行更简单的批处理数据并行性。

• 创建不一定由网格和分区规范组成的 Sharding 的方法。如果您愿意，可以充分利用 OpSharding 的灵活性或您想要的任何其他 Sharding。

• 还有更多

示例

import jax
import jax.numpy as jnp
from jax.sharding import PartitionSpec as P
import numpy as np
x = jnp.arange(8)

# Let's say there are 8 devices in jax.devices()
mesh = jax.sharding.Mesh(np.array(jax.devices()).reshape(4, 2), ('x', 'y'))
sharding = jax.sharding.NamedSharding(mesh, P('x'))

sharded_x = jax.device_put(x, sharding)

# `matmul_sharded_x` and `sin_sharded_x` are sharded. `jit` is able to operate over a
# sharded array without copying data to a single device.
matmul_sharded_x = sharded_x @ sharded_x.T
sin_sharded_x = jnp.sin(sharded_x)

# Even jnp.copy preserves the sharding on the output.
copy_sharded_x = jnp.copy(sharded_x)

# double_out is also sharded
double_out = jax.jit(lambda x: x * 2)(sharded_x)

启用 jax.Array 后可能出现哪些问题？

名为 jax.Array 的新公共类型

所有 isinstance(..., jnp.DeviceArray) 或 isinstance(.., jax.xla.DeviceArray) 以及其他 DeviceArray 的变体都应切换为使用 isinstance(..., jax.Array)。

由于 jax.Array 可以表示 DA、SDA 和 GDA，您可以通过以下方式区分 jax.Array 中的这 3 种类型：

• x.is_fully_addressable and len(x.sharding.device_set) == 1 – 这意味着 jax.Array 类似于 DA

• x.is_fully_addressable and (len(x.sharding.device_set) > 1 – 这意味着 jax.Array 类似于 SDA

• not x.is_fully_addressable – 这意味着 jax.Array 类似于 GDA，并跨多个进程

对于 ShardedDeviceArray，您可以将 isinstance(..., pxla.ShardedDeviceArray) 移至 isinstance(..., jax.Array) and x.is_fully_addressable and len(x.sharding.device_set) > 1。

一般来说，无法将 1 个设备上的 ShardedDeviceArray 与任何其他类型的单设备 Array 区分开来。

GDA 的 API 名称更改

GDA 的 local_shards 和 local_data 已被弃用。

请使用与 jax.Array 和 GDA 兼容的 addressable_shards 和 addressable_data。

创建 jax.Array

当 jax_array 标志为 True 时，所有 JAX 函数都将输出 jax.Array。如果您之前使用 GlobalDeviceArray.from_callback 或 make_sharded_device_array 或 make_device_array 函数来显式创建相应的 JAX 数据类型，则需要将它们切换为使用 jax.make_array_from_callback() 或 jax.make_array_from_single_device_arrays()。

对于 GDA

GlobalDeviceArray.from_callback(shape, mesh, pspec, callback) 可以 1:1 切换为 jax.make_array_from_callback(shape, jax.sharding.NamedSharding(mesh, pspec), callback)。

如果您之前使用原始 GDA 构造函数创建 GDA，则执行以下操作

GlobalDeviceArray(shape, mesh, pspec, buffers) 可以变为 jax.make_array_from_single_device_arrays(shape, jax.sharding.NamedSharding(mesh, pspec), buffers)

对于 SDA

make_sharded_device_array(aval, sharding_spec, device_buffers, indices) 可以变为 jax.make_array_from_single_device_arrays(shape, sharding, device_buffers)。

要决定 sharding 应该是什么，取决于您创建 SDA 的原因

如果创建它是为了作为 pmap 的输入，则 sharding 可以是：jax.sharding.PmapSharding(devices, sharding_spec)。

如果创建它是为了作为 pjit 的输入，则 sharding 可以是 jax.sharding.NamedSharding(mesh, pspec)。

切换到使用 jax.Array 作为主机本地输入后，pjit 的重大变更

如果您仅使用 GDA 参数调用 pjit，则可以跳过本节！🎉

启用 jax.Array 后，pjit 的所有输入都必须是全局形状的。这是一个与之前行为的重大变更，之前的行为中 pjit 会将进程本地参数连接成一个全局值；此连接操作不再发生。

为什么要进行此重大变更？现在，每个数组都明确说明了其本地分片如何融入一个全局整体，而不是隐式地说明。这种更明确的表示形式也解锁了额外的灵活性，例如，在某些 TPU 模型上使用具有 pjit 的非连续网格可以提高效率。

当启用 jax.Array 时，运行多进程 pjit 计算并传递主机本地输入可能会导致类似于此的错误

示例

网格 = {'x': 2, 'y': 2, 'z': 2} 并且主机本地输入形状 == (4,) 且 pspec = P(('x', 'y', 'z'))

由于 pjit 不会使用 jax.Array 将主机本地形状提升为全局形状，因此您会收到以下错误

注意：只有当您的主机本地形状小于网格的形状时，您才会看到此错误。

ValueError: One of pjit arguments was given the sharding of
NamedSharding(mesh={'x': 2, 'y': 2, 'chips': 2}, partition_spec=PartitionSpec(('x', 'y', 'chips'),)),
which implies that the global size of its dimension 0 should be divisible by 8,
but it is equal to 4

这个错误是有道理的，因为当维度 0 上的值为 4 时，您无法将维度 0 分片为 8 个部分。

如果您仍然将主机本地输入传递给 pjit，您该如何迁移？我们正在提供过渡 API 来帮助您迁移

注意：如果您在单个进程上运行 pjitted 计算，则不需要这些实用程序。

from jax.experimental import multihost_utils

global_inps = multihost_utils.host_local_array_to_global_array(
    local_inputs, mesh, in_pspecs)

global_outputs = pjit(f, in_shardings=in_pspecs,
                      out_shardings=out_pspecs)(global_inps)

local_outs = multihost_utils.global_array_to_host_local_array(
    global_outputs, mesh, out_pspecs)

host_local_array_to_global_array 是一种类型转换，它会查看仅具有本地分片的值，并将其本地形状更改为 pjit 在更改之前传递该值时会假定的形状。

仍然支持传递完全复制的输入，即每个进程上具有相同形状，并且 in_axis_resources 为 P(None) 的情况。在这种情况下，您不必使用 host_local_array_to_global_array，因为形状已经是全局的。

key = jax.random.PRNGKey(1)

# As you can see, using host_local_array_to_global_array is not required since in_axis_resources says
# that the input is fully replicated via P(None)
pjit(f, in_shardings=None, out_shardings=None)(key)

# Mixing inputs
global_inp = multihost_utils.host_local_array_to_global_array(
    local_inp, mesh, P('data'))
global_out = pjit(f, in_shardings=(P(None), P('data')),
                  out_shardings=...)(key, global_inp)

FROM_GDA 和 jax.Array

如果您在 pjit 的 in_axis_resources 参数中使用了 FROM_GDA，那么使用 jax.Array 后，无需向 in_axis_resources 传递任何内容，因为 jax.Array 将遵循计算跟随分片的语义。

例如

pjit(f, in_shardings=FROM_GDA, out_shardings=...) can be replaced by pjit(f, out_shardings=...)

如果您的输入（如 numpy 数组等）中混合使用了 PartitionSpecs 和 FROM_GDA，请使用 host_local_array_to_global_array 将它们转换为 jax.Array。

例如

如果您有以下代码

pjitted_f = pjit(
    f, in_shardings=(FROM_GDA, P('x'), FROM_GDA, P(None)),
    out_shardings=...)
pjitted_f(gda1, np_array1, gda2, np_array2)

那么您可以将其替换为

pjitted_f = pjit(f, out_shardings=...)

array2, array3 = multihost_utils.host_local_array_to_global_array(
    (np_array1, np_array2), mesh, (P('x'), P(None)))

pjitted_f(array1, array2, array3, array4)

live_buffers 替换为 live_arrays

jax Device 上的 live_buffers 属性已弃用。请改用与 jax.Array 兼容的 jax.live_arrays()。

处理类似批次的主机本地输入到 pjit

如果您在多进程环境中将主机本地输入传递给 pjit，请使用 multihost_utils.host_local_array_to_global_array 将批次转换为全局 jax.Array，然后将其传递给 pjit。

此类主机本地输入最常见的示例是输入数据批次。

这将适用于任何主机本地输入（不仅仅是输入数据批次）。

from jax.experimental import multihost_utils

batch = multihost_utils.host_local_array_to_global_array(
    batch, mesh, batch_partition_spec)

有关此更改和更多示例，请参阅上面的 pjit 部分。

RecursionError：递归调用 jit

当您代码的某些部分禁用 jax.Array，然后您仅在其他部分启用它时，会发生这种情况。例如，如果您使用一些第三方代码，该代码禁用了 jax.Array，并且您从该库中获取了一个 DeviceArray，然后在您的库中启用 jax.Array 并将该 DeviceArray 传递给 JAX 函数，则会导致 RecursionError。

当默认启用 jax.Array 以便所有库都返回 jax.Array（除非它们明确禁用它）时，此错误应消失。