分布式数组和自动并行化

本教程讨论了通过 jax.Array 实现的并行化，这是 JAX v0.4.1 及更高版本中提供的统一数组对象模型。有关使用此 API 的真实机器学习训练示例，请参阅 The Training Cookbook。

from typing import Optional

import numpy as np

import jax
import jax.numpy as jnp

⚠️ 警告：此笔记本需要 8 个设备才能运行。

if len(jax.local_devices()) < 8:
  raise Exception("Notebook requires 8 devices to run")

简介和快速示例

通过阅读本教程笔记本，您将了解 jax.Array，这是一种统一的数据类型，用于表示数组，即使物理存储跨越多个设备。您还将了解如何将 jax.Array 与 jax.jit 一起使用，以提供基于编译器的自动并行化。

在我们分步讲解之前，这里有一个快速示例。首先，我们将创建一个跨多个设备分片的 jax.Array。

from jax.sharding import PartitionSpec as P, NamedSharding

# Create a Sharding object to distribute a value across devices:
mesh = jax.make_mesh((4, 2), ('x', 'y'))

# Create an array of random values:
x = jax.random.normal(jax.random.key(0), (8192, 8192))
# and use jax.device_put to distribute it across devices:
y = jax.device_put(x, NamedSharding(mesh, P('x', 'y')))
jax.debug.visualize_array_sharding(y)

┌──────────┬──────────┐
│  TPU 0   │  TPU 1   │
├──────────┼──────────┤
│  TPU 2   │  TPU 3   │
├──────────┼──────────┤
│  TPU 6   │  TPU 7   │
├──────────┼──────────┤
│  TPU 4   │  TPU 5   │
└──────────┴──────────┘

接下来，我们将对其应用计算，并可视化结果值如何在多个设备上存储。

z = jnp.sin(y)
jax.debug.visualize_array_sharding(z)

┌──────────┬──────────┐
│  TPU 0   │  TPU 1   │
├──────────┼──────────┤
│  TPU 2   │  TPU 3   │
├──────────┼──────────┤
│  TPU 6   │  TPU 7   │
├──────────┼──────────┤
│  TPU 4   │  TPU 5   │
└──────────┴──────────┘

jnp.sin 应用的评估已自动并行化到存储输入值（和输出值）的设备上。

# `x` is present on a single device
%timeit -n 5 -r 5 jnp.sin(x).block_until_ready()

The slowest run took 8.96 times longer than the fastest. This could mean that an intermediate result is being cached.
25.2 ms ± 30.9 ms per loop (mean ± std. dev. of 5 runs, 5 loops each)

# `y` is sharded across 8 devices.
%timeit -n 5 -r 5 jnp.sin(y).block_until_ready()

2.4 ms ± 61.4 µs per loop (mean ± std. dev. of 5 runs, 5 loops each)

现在让我们更详细地了解这些部分！

Sharding 描述了数组值如何在内存中跨设备布局

Sharding 基础知识，以及 NamedSharding 子类

要在多个设备上并行计算，我们首先必须将输入数据布局到多个设备上。

在 JAX 中，Sharding 对象描述分布式内存布局。它们可以与 jax.device_put 一起使用，以生成具有分布式布局的值。

例如，这里有一个具有单设备 Sharding 的值。

import jax
x = jax.random.normal(jax.random.key(0), (8192, 8192))

jax.debug.visualize_array_sharding(x)

┌───────────────────────┐
│                       │
│                       │
│                       │
│                       │
│         TPU 0         │
│                       │
│                       │
│                       │
│                       │
└───────────────────────┘

在这里，我们使用 jax.debug.visualize_array_sharding 函数来显示值 x 存储在内存中的位置。所有 x 都存储在单个设备上，因此可视化非常无聊！

但是，我们可以通过使用 jax.device_put 和 Sharding 对象将 x 分片到多个设备上。首先，我们使用 jax.make_mesh 创建一个 Devices 的 numpy.ndarray，它会考虑 Device 顺序的硬件拓扑。

from jax.sharding import Mesh, PartitionSpec, NamedSharding

P = PartitionSpec

mesh = jax.make_mesh((4, 2), ('a', 'b'))
y = jax.device_put(x, NamedSharding(mesh, P('a', 'b')))
jax.debug.visualize_array_sharding(y)

┌──────────┬──────────┐
│  TPU 0   │  TPU 1   │
├──────────┼──────────┤
│  TPU 2   │  TPU 3   │
├──────────┼──────────┤
│  TPU 6   │  TPU 7   │
├──────────┼──────────┤
│  TPU 4   │  TPU 5   │
└──────────┴──────────┘

我们可以定义一个辅助函数来简化事情。

default_mesh = jax.make_mesh((4, 2), ('a', 'b'))

def mesh_sharding(
    pspec: PartitionSpec, mesh: Optional[Mesh] = None,
  ) -> NamedSharding:
  if mesh is None:
    mesh = default_mesh
  return NamedSharding(mesh, pspec)

y = jax.device_put(x, mesh_sharding(P('a', 'b')))
jax.debug.visualize_array_sharding(y)

┌──────────┬──────────┐
│  TPU 0   │  TPU 1   │
├──────────┼──────────┤
│  TPU 2   │  TPU 3   │
├──────────┼──────────┤
│  TPU 6   │  TPU 7   │
├──────────┼──────────┤
│  TPU 4   │  TPU 5   │
└──────────┴──────────┘

在这里，我们使用 P('a', 'b') 来表示 x 的第一个和第二个轴应该分别在设备网格轴 'a' 和 'b' 上分片。我们可以轻松切换到 P('b', 'a') 将 x 的轴分片到不同的设备上。

y = jax.device_put(x, mesh_sharding(P('b', 'a')))
jax.debug.visualize_array_sharding(y)

┌───────┬───────┬───────┬───────┐
│       │       │       │       │
│ TPU 0 │ TPU 2 │ TPU 6 │ TPU 4 │
│       │       │       │       │
│       │       │       │       │
├───────┼───────┼───────┼───────┤
│       │       │       │       │
│ TPU 1 │ TPU 3 │ TPU 7 │ TPU 5 │
│       │       │       │       │
│       │       │       │       │
└───────┴───────┴───────┴───────┘

# This `None` means that `x` is not sharded on its second dimension,
# and since the Mesh axis name 'b' is not mentioned, shards are
# replicated across it.
y = jax.device_put(x, mesh_sharding(P('a', None)))
jax.debug.visualize_array_sharding(y)

┌───────────────────────┐
│        TPU 0,1        │
├───────────────────────┤
│        TPU 2,3        │
├───────────────────────┤
│        TPU 6,7        │
├───────────────────────┤
│        TPU 4,5        │
└───────────────────────┘

在这里，因为 P('a', None) 没有提到 Mesh 轴名称 'b'，所以我们在 'b' 轴上获得了复制。这里的 None 只是作为一个占位符，与 x 的第二个轴对齐，而不表达在任何网格轴上的分片。（作为简写，可以省略尾随的 None，因此 P('a', None) 与 P('a') 含义相同。但明确一点并没有坏处！）

要仅在 x 的第二个轴上分片，我们可以在 PartitionSpec 中使用 None 占位符。

y = jax.device_put(x, mesh_sharding(P(None, 'b')))
jax.debug.visualize_array_sharding(y)

┌───────────┬───────────┐
│           │           │
│           │           │
│           │           │
│           │           │
│TPU 0,2,4,6│TPU 1,3,5,7│
│           │           │
│           │           │
│           │           │
│           │           │
└───────────┴───────────┘

y = jax.device_put(x, mesh_sharding(P(None, 'a')))
jax.debug.visualize_array_sharding(y)

┌───────┬───────┬───────┬───────┐
│       │       │       │       │
│       │       │       │       │
│       │       │       │       │
│       │       │       │       │
│TPU 0,1│TPU 2,3│TPU 6,7│TPU 4,5│
│       │       │       │       │
│       │       │       │       │
│       │       │       │       │
│       │       │       │       │
└───────┴───────┴───────┴───────┘

对于固定的网格，我们甚至可以将 x 的一个逻辑轴分片到多个设备网格轴上。

y = jax.device_put(x, mesh_sharding(P(('a', 'b'), None)))
jax.debug.visualize_array_sharding(y)

┌───────────────────────┐
│         TPU 0         │
├───────────────────────┤
│         TPU 1         │
├───────────────────────┤
│         TPU 2         │
├───────────────────────┤
│         TPU 3         │
├───────────────────────┤
│         TPU 6         │
├───────────────────────┤
│         TPU 7         │
├───────────────────────┤
│         TPU 4         │
├───────────────────────┤
│         TPU 5         │
└───────────────────────┘

使用 NamedSharding 可以方便地一次定义一个设备网格并为其轴命名，然后在每个 device_put 的 PartitionSpec 中按需引用这些名称。

计算遵循数据分片并自动并行化

通过分片输入数据，编译器可以为我们提供并行计算。特别是，用 jax.jit 装饰的函数可以在不将数据复制到单个设备的情况下操作分片数组。相反，计算遵循分片：基于输入数据的分片，编译器会决定中间值和输出值的分片，并并行化它们的评估，甚至在必要时插入通信操作。

例如，最简单的计算是逐元素计算。

mesh = jax.make_mesh((4, 2), ('a', 'b'))

x = jax.device_put(x, NamedSharding(mesh, P('a', 'b')))
print('input sharding:')
jax.debug.visualize_array_sharding(x)

y = jnp.sin(x)
print('output sharding:')
jax.debug.visualize_array_sharding(y)

input sharding:
output sharding:

┌──────────┬──────────┐
│  TPU 0   │  TPU 1   │
├──────────┼──────────┤
│  TPU 2   │  TPU 3   │
├──────────┼──────────┤
│  TPU 6   │  TPU 7   │
├──────────┼──────────┤
│  TPU 4   │  TPU 5   │
└──────────┴──────────┘

┌──────────┬──────────┐
│  TPU 0   │  TPU 1   │
├──────────┼──────────┤
│  TPU 2   │  TPU 3   │
├──────────┼──────────┤
│  TPU 6   │  TPU 7   │
├──────────┼──────────┤
│  TPU 4   │  TPU 5   │
└──────────┴──────────┘

在这里，对于逐元素操作 jnp.sin，编译器选择的输出分片与输入相同。此外，编译器自动并行化了计算，因此每个设备并行地计算其输入分片上的输出分片。

换句话说，即使我们编写 jnp.sin 计算的方式就像单台机器执行它一样，编译器也会为我们拆分计算并在多个设备上执行。

我们也可以对非逐元素操作进行同样的操作。考虑一个具有分片输入的矩阵乘法。

y = jax.device_put(x, NamedSharding(mesh, P('a', None)))
z = jax.device_put(x, NamedSharding(mesh, P(None, 'b')))
print('lhs sharding:')
jax.debug.visualize_array_sharding(y)
print('rhs sharding:')
jax.debug.visualize_array_sharding(z)

w = jnp.dot(y, z)
print('out sharding:')
jax.debug.visualize_array_sharding(w)

lhs sharding:
rhs sharding:
out sharding:

┌───────────────────────┐
│        TPU 0,1        │
├───────────────────────┤
│        TPU 2,3        │
├───────────────────────┤
│        TPU 6,7        │
├───────────────────────┤
│        TPU 4,5        │
└───────────────────────┘

┌───────────┬───────────┐
│           │           │
│           │           │
│           │           │
│           │           │
│TPU 0,2,4,6│TPU 1,3,5,7│
│           │           │
│           │           │
│           │           │
│           │           │
└───────────┴───────────┘

┌──────────┬──────────┐
│  TPU 0   │  TPU 1   │
├──────────┼──────────┤
│  TPU 2   │  TPU 3   │
├──────────┼──────────┤
│  TPU 6   │  TPU 7   │
├──────────┼──────────┤
│  TPU 4   │  TPU 5   │
└──────────┴──────────┘

在这里，编译器选择的输出分片可以最大限度地并行化计算：无需通信，每个设备已经拥有计算其输出分片所需的输入分片。

我们如何确定它实际上是在并行运行的？我们可以做一个简单的计时实验。

x_single = jax.device_put(x, jax.devices()[0])
jax.debug.visualize_array_sharding(x_single)

┌───────────────────────┐
│                       │
│                       │
│                       │
│                       │
│         TPU 0         │
│                       │
│                       │
│                       │
│                       │
└───────────────────────┘

np.allclose(jnp.dot(x_single, x_single),
            jnp.dot(y, z))

True

%timeit -n 5 -r 5 jnp.dot(x_single, x_single).block_until_ready()

49.7 ms ± 349 µs per loop (mean ± std. dev. of 5 runs, 5 loops each)

%timeit -n 5 -r 5 jnp.dot(y, z).block_until_ready()

7.47 ms ± 44.8 µs per loop (mean ± std. dev. of 5 runs, 5 loops each)

即使复制一个分片的 Array 也会产生一个具有输入分片的分片结果。

w_copy = jnp.copy(w)
jax.debug.visualize_array_sharding(w_copy)

┌──────────┬──────────┐
│  TPU 0   │  TPU 1   │
├──────────┼──────────┤
│  TPU 2   │  TPU 3   │
├──────────┼──────────┤
│  TPU 6   │  TPU 7   │
├──────────┼──────────┤
│  TPU 4   │  TPU 5   │
└──────────┴──────────┘

因此，计算遵循数据放置：当我们使用 jax.device_put 显式分片数据，并对该数据应用函数时，编译器会尝试并行化计算并决定输出分片。这种分片数据的策略是 JAX 的显式设备放置策略 的推广。

当显式分片不一致时，JAX 会报错

但是，如果计算的两个参数被显式放置在不同的设备集上，或者设备顺序不兼容怎么办？在这些模糊的情况下，会引发错误。

import textwrap
from termcolor import colored

def print_exception(e):
  name = colored(f'{type(e).__name__}', 'red', force_color=True)
  print(textwrap.fill(f'{name}: {str(e)}'))

sharding1 = NamedSharding(Mesh(jax.devices()[:4], 'x'), P('x'))
sharding2 = NamedSharding(Mesh(jax.devices()[4:], 'x'), P('x'))

y = jax.device_put(x, sharding1)
z = jax.device_put(x, sharding2)
try: y + z
except ValueError as e: print_exception(e)

ValueError: Received incompatible devices for jitted
computation. Got argument x1 of jax.numpy.add with shape int32[24] and
device ids [0, 1, 2, 3] on platform TPU and argument x2 of
jax.numpy.add with shape int32[24] and device ids [4, 5, 6, 7] on
platform TPU

devices = jax.devices()
permuted_devices = [devices[i] for i in [0, 1, 2, 3, 6, 7, 4, 5]]

sharding1 = NamedSharding(Mesh(devices, 'x'), P('x'))
sharding2 = NamedSharding(Mesh(permuted_devices, 'x'), P('x'))

y = jax.device_put(x, sharding1)
z = jax.device_put(x, sharding2)
try: y + z
except ValueError as e: print_exception(e)

ValueError: Received incompatible devices for jitted
computation. Got argument x1 of jax.numpy.add with shape int32[24] and
device ids [0, 1, 2, 3, 4, 5, 6, 7] on platform TPU and argument x2 of
jax.numpy.add with shape int32[24] and device ids [0, 1, 2, 3, 6, 7,
4, 5] on platform TPU

我们说使用 jax.device_put 显式放置或分片的数组会被 *承诺* 到它们的设备上，因此不会被自动移动。有关更多信息，请参阅 设备放置 FAQ。

当数组 *未* 被显式放置或使用 jax.device_put 分片时，它们会被 *未承诺* 地放置在默认设备上。与已承诺数组不同，未承诺数组可以被自动移动和重新分片：也就是说，未承诺数组可以作为计算的参数，即使其他参数被显式放置在不同的设备上。

例如，jnp.zeros、jnp.arange 和 jnp.array 的输出是未承诺的。

y = jax.device_put(x, sharding1)
y + jnp.ones_like(y)
y + jnp.arange(y.size).reshape(y.shape)
print('no error!')

no error!

在 jit 编译的代码中约束中间值的分片

虽然编译器会尝试决定函数的中间值和输出应该如何分片，但我们也可以使用 jax.lax.with_sharding_constraint 给它一些提示。使用 jax.lax.with_sharding_constraint 类似于 jax.device_put，除了我们在已编译（即使用 jit 装饰）的函数内部使用它。

mesh = jax.make_mesh((4, 2), ('x', 'y'))

x = jax.random.normal(jax.random.key(0), (8192, 8192))
x = jax.device_put(x, NamedSharding(mesh, P('x', 'y')))

@jax.jit
def f(x):
  x = x + 1
  y = jax.lax.with_sharding_constraint(x, NamedSharding(mesh, P('y', 'x')))
  return y

jax.debug.visualize_array_sharding(x)
y = f(x)
jax.debug.visualize_array_sharding(y)

┌──────────┬──────────┐
│  TPU 0   │  TPU 1   │
├──────────┼──────────┤
│  TPU 2   │  TPU 3   │
├──────────┼──────────┤
│  TPU 6   │  TPU 7   │
├──────────┼──────────┤
│  TPU 4   │  TPU 5   │
└──────────┴──────────┘

┌───────┬───────┬───────┬───────┐
│       │       │       │       │
│ TPU 0 │ TPU 2 │ TPU 6 │ TPU 4 │
│       │       │       │       │
│       │       │       │       │
├───────┼───────┼───────┼───────┤
│       │       │       │       │
│ TPU 1 │ TPU 3 │ TPU 7 │ TPU 5 │
│       │       │       │       │
│       │       │       │       │
└───────┴───────┴───────┴───────┘

@jax.jit
def f(x):
  x = x + 1
  y = jax.lax.with_sharding_constraint(x, NamedSharding(mesh, P()))
  return y

jax.debug.visualize_array_sharding(x)
y = f(x)
jax.debug.visualize_array_sharding(y)

┌──────────┬──────────┐
│  TPU 0   │  TPU 1   │
├──────────┼──────────┤
│  TPU 2   │  TPU 3   │
├──────────┼──────────┤
│  TPU 6   │  TPU 7   │
├──────────┼──────────┤
│  TPU 4   │  TPU 5   │
└──────────┴──────────┘

┌───────────────────────┐
│                       │
│                       │
│                       │
│                       │
│  TPU 0,1,2,3,4,5,6,7  │
│                       │
│                       │
│                       │
│                       │
└───────────────────────┘

通过添加 with_sharding_constraint，我们约束了输出的分片。除了遵循特定中间值的注解外，编译器还将使用注解来决定其他值的分片。

通常，根据值最终如何被使用，为计算的输注解分片是一个好习惯。

示例：神经网络

⚠️ 警告：以下内容旨在简单演示 jax.Array 的自动分片传播，但可能不反映实际示例的最佳实践。 例如，实际示例可能需要更多地使用 with_sharding_constraint。

我们可以使用 jax.device_put 和 jax.jit 的计算遵循分片特性来并行化神经网络中的计算。这里有一些简单的例子，基于这个基本的神经网络。

import jax
import jax.numpy as jnp

def predict(params, inputs):
  for W, b in params:
    outputs = jnp.dot(inputs, W) + b
    inputs = jnp.maximum(outputs, 0)
  return outputs

def loss(params, batch):
  inputs, targets = batch
  predictions = predict(params, inputs)
  return jnp.mean(jnp.sum((predictions - targets)**2, axis=-1))

loss_jit = jax.jit(loss)
gradfun = jax.jit(jax.grad(loss))

def init_layer(key, n_in, n_out):
  k1, k2 = jax.random.split(key)
  W = jax.random.normal(k1, (n_in, n_out)) / jnp.sqrt(n_in)
  b = jax.random.normal(k2, (n_out,))
  return W, b

def init_model(key, layer_sizes, batch_size):
  key, *keys = jax.random.split(key, len(layer_sizes))
  params = list(map(init_layer, keys, layer_sizes[:-1], layer_sizes[1:]))

  key, *keys = jax.random.split(key, 3)
  inputs = jax.random.normal(keys[0], (batch_size, layer_sizes[0]))
  targets = jax.random.normal(keys[1], (batch_size, layer_sizes[-1]))

  return params, (inputs, targets)

layer_sizes = [784, 8192, 8192, 8192, 10]
batch_size = 8192

params, batch = init_model(jax.random.key(0), layer_sizes, batch_size)

8 路批数据并行

mesh = jax.make_mesh((8,), ('batch',))

sharding = NamedSharding(mesh, P('batch'))
replicated_sharding = NamedSharding(mesh, P())

batch = jax.device_put(batch, sharding)
params = jax.device_put(params, replicated_sharding)

loss_jit(params, batch)

Array(23.469475, dtype=float32)

step_size = 1e-5

for _ in range(30):
  grads = gradfun(params, batch)
  params = [(W - step_size * dW, b - step_size * db)
            for (W, b), (dW, db) in zip(params, grads)]

print(loss_jit(params, batch))

10.760109

%timeit -n 5 -r 5 gradfun(params, batch)[0][0].block_until_ready()

53.8 ms ± 1.14 ms per loop (mean ± std. dev. of 5 runs, 5 loops each)

batch_single = jax.device_put(batch, jax.devices()[0])
params_single = jax.device_put(params, jax.devices()[0])

%timeit -n 5 -r 5 gradfun(params_single, batch_single)[0][0].block_until_ready()

351 ms ± 81.2 ms per loop (mean ± std. dev. of 5 runs, 5 loops each)

4 路批数据并行和 2 路模型张量并行

mesh = jax.make_mesh((4, 2), ('batch', 'model'))

batch = jax.device_put(batch, NamedSharding(mesh, P('batch', None)))
jax.debug.visualize_array_sharding(batch[0])
jax.debug.visualize_array_sharding(batch[1])

┌───────┐
│TPU 0,1│
├───────┤
│TPU 2,3│
├───────┤
│TPU 6,7│
├───────┤
│TPU 4,5│
└───────┘

┌───────┐
│TPU 0,1│
├───────┤
│TPU 2,3│
├───────┤
│TPU 6,7│
├───────┤
│TPU 4,5│
└───────┘

replicated_sharding = NamedSharding(mesh, P())

(W1, b1), (W2, b2), (W3, b3), (W4, b4) = params

W1 = jax.device_put(W1, replicated_sharding)
b1 = jax.device_put(b1, replicated_sharding)

W2 = jax.device_put(W2, NamedSharding(mesh, P(None, 'model')))
b2 = jax.device_put(b2, NamedSharding(mesh, P('model')))

W3 = jax.device_put(W3, NamedSharding(mesh, P('model', None)))
b3 = jax.device_put(b3, replicated_sharding)

W4 = jax.device_put(W4, replicated_sharding)
b4 = jax.device_put(b4, replicated_sharding)

params = (W1, b1), (W2, b2), (W3, b3), (W4, b4)

jax.debug.visualize_array_sharding(W2)

┌───────────┬───────────┐
│           │           │
│           │           │
│           │           │
│           │           │
│TPU 0,2,4,6│TPU 1,3,5,7│
│           │           │
│           │           │
│           │           │
│           │           │
└───────────┴───────────┘

jax.debug.visualize_array_sharding(W3)

┌───────────────────────┐
│                       │
│      TPU 0,2,4,6      │
│                       │
│                       │
├───────────────────────┤
│                       │
│      TPU 1,3,5,7      │
│                       │
│                       │
└───────────────────────┘

print(loss_jit(params, batch))

10.760109

step_size = 1e-5

for _ in range(30):
    grads = gradfun(params, batch)
    params = [(W - step_size * dW, b - step_size * db)
              for (W, b), (dW, db) in zip(params, grads)]

print(loss_jit(params, batch))

10.752513

(W1, b1), (W2, b2), (W3, b3), (W4, b4) = params
jax.debug.visualize_array_sharding(W2)
jax.debug.visualize_array_sharding(W3)

┌───────────┬───────────┐
│           │           │
│           │           │
│           │           │
│           │           │
│TPU 0,2,4,6│TPU 1,3,5,7│
│           │           │
│           │           │
│           │           │
│           │           │
└───────────┴───────────┘

┌───────────────────────┐
│                       │
│      TPU 0,2,4,6      │
│                       │
│                       │
├───────────────────────┤
│                       │
│      TPU 1,3,5,7      │
│                       │
│                       │
└───────────────────────┘

%timeit -n 10 -r 10 gradfun(params, batch)[0][0].block_until_ready()

51.4 ms ± 454 µs per loop (mean ± std. dev. of 10 runs, 10 loops each)