jax.Array 迁移

yashkatariya@

太长不看

自版本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 设置为类似真值的内容（例如，1）；

• 如果你的代码使用 absl 解析标志，请将布尔标志 jax_array 设置为类似真值的内容。

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

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

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

判断 jax.Array 是否是问题根源的最简单方法是禁用 jax.Array 并查看问题是否消失。

我该如何暂时禁用 jax.Array？

截至 2023年3月15日 ，可以通过以下方式禁用 jax.Array：

• 将 shell 环境变量 JAX_ARRAY 设置为假值（例如，0）；

• 如果你的代码使用 absl 解析标志，请将布尔标志 jax_array 设置为假值。

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

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

为什么要创建 jax.Array？

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

我们还引入了一个新的 Sharding 抽象，描述了逻辑数组如何在物理上分布在一个或多个设备上，如 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 中区分这三种类型：

• 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) 并且 x.is_fully_addressable 并且 len(x.sharding.device_set) > 1。

通常情况下，无法区分一个在1个设备上的 ShardedDeviceArray 与其他任何类型的单设备数组。

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) 可以一对一地转换为 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)。

要决定分片应该是什么，这取决于你创建SDAs的原因：

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

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

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

如果你只使用 GDA 参数来 pjit，你可以跳过这一部分！🎉

启用 jax.Array 后，所有输入到 pjit 的参数必须是全局形状的。这与之前的行为有所不同，之前 pjit 会将进程本地参数连接成一个全局值；这种连接不再发生。

为什么我们要进行这个破坏性的更改？现在每个数组都明确说明了其本地分片如何适应全局整体，而不是隐含地处理。更明确的表示还解锁了额外的灵活性，例如在使用 pjit 时可以使用非连续网格，这可以提高某些TPU模型的效率。

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

示例：

Mesh = {'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 时，将其分片为 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 在更改之前如果传递该值时会假设的形状。

传递完全复制的输入，即在每个进程上具有相同形状，并将 P(None) 作为 in_axis_resources 仍然受支持。在这种情况下，您不需要使用 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.live_arrays()，它与 jax.Array 兼容。

处理主机本地输入到 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 函数，这将导致一个递归错误。

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