Python 是数据科学、机器学习和数值计算领域最常用的编程语言。它在科学家和研究人员中日益受欢迎。在 Python 生态系统中， NumPy 是执行基于数组的数值计算的基础 Python 库。
NumPy 的标准实现可在单个 CPU 核心上运行，只有一些操作可以跨核心并行执行。这种单线程、仅使用 CPU 的执行方式限制了可处理的数据规模，也限制了执行计算的速度。
虽然可以使用 GPU 加速的 NumPy 实现，但跨多个 GPU 或节点扩展基于 NumPy 的代码通常需要大量的代码修改，包括手动数据分区和同步以及用于分布式执行的数据移动。这种代码更改可能十分复杂且耗时，以确保功能正确且性能出色。
此外，在分布式编程方面缺乏专业知识的领域科学家通常会与计算机科学专家合作或咨询，以完成更改，这进一步减缓了实验和验证研究的过程。
为解决这一生产力问题，我们构建了 NVIDIA cuPyNumeric ，这是 NumPy API 的开源、分布式和加速实现，旨在直接替代 NumPy。借助 cuPyNumeric，科学家、研究人员和处理大规模问题的人员可以使用笔记本电脑或台式机在中等大小的数据集上开发和测试 NumPy 程序。然后，他们可以使用超级计算机或云自行扩展相同的程序，生成大型数据集，而无需更改代码。
cuPyNumeric 为您提供 NumPy 的工作效率，以及加速分布式 GPU 计算的性能优势。cuPyNumeric 支持基本的 NumPy 功能，包括就地更新、 广播 和 综合索引语义 。
cuPyNumeric：直接替代 NumPy
您可以通过在 NumPy 程序中导入 cupynumeric 而不是 numpy 来开始使用 cuPyNumeric。以下代码示例使用 NumPy API 近似表示π：
# import numpy as np
import cupynumeric as np
num_points = 1000000000
# Generates random points in a square of area 4
x = np.random.uniform(-1, 1, size=num_points)
y = np.random.uniform(-1, 1, size=num_points)
# Calculates the probability of points falling in a circle of radius 1
probability_in_circle = np.sum((x ** 2 + y ** 2) < 1) / num_points
# probability_in_circle = (area of the circle) / (area of the square)
# ==> PI = probability_in_circle * (area of the square)
print(f"PI is {probability_in_circle * 4}")
无论使用 cuPyNumeric 还是 NumPy 运行，该程序都会生成以下输出，尽管不同运行的确切值可能不同：
PI is 3.141572908
只需对 import 语句进行简单更改，此代码即可扩展到任何机器：
- 在没有 GPU 的笔记本电脑上，代码使用可用的 CPU 核心进行并行化。
- 在具有多个 GPU 的系统上，由于多个 GPU 的加速，相同的代码运行速度要快得多。
- 尽管代码没有足够的计算能力来充分利用如此强大的机器，但这些代码甚至可以在不修改代码的情况下在超级计算机上运行。
cuPyNumeric 中的模板示例
这是一个更有趣的示例，展示了 cuPyNumeric 的优势。模板计算是科学计算中最常见的算法类型之一。模板程序在 NumPy 中自然使用将相邻单元与中心单元移动和对齐的切片来表示。
但是，对于多 GPU 系统，这些程序并不总是简单的并行化，因为当数组在多个 GPU 上进行分区时，一个 GPU 对中心单元的任何更改都必须传播到其他 GPU。
cuPyNumeric 可透明地将使用纯 NumPy 编写的模板代码扩展到任意数量的 GPU 或节点，而无需程序员考虑此分布。
以下代码示例创建了 2D 数组网格，并通过引用数组的别名切片对每个单元执行 5 点 stencil 计算：
import cupynumeric as np
grid = np.random.rand(N+2, N+2)
# Create multiple aliasing views of the grid array.
center = grid[1:-1, 1:-1]
north = grid[0:-2, 1:-1]
east = grid[1:-1, 2: ]
west = grid[1:-1, 0:-2]
south = grid[2: , 1:-1]
for _ in range(niters):
average = (center + north + east + west + south) * 0.2
center[:] = average
当此程序在多个 GPU 上运行时，系统会自动将数组分割成机器的 GPU 上的图块，并分配运算，以便每个 GPU 对其本地的数据块执行运算。
除数据和计算分区外，cuPyNumeric 还会自动推断网格阵列的别名视图之间的通信。在操作 center[:] = average 之后，中心数组的更新必须传播到 north、east、west 和 south 数组。
cuPyNumeric 在程序中无需任何显式通信代码即可执行此通信，从而自动发现节点内或跨节点的 GPU 之间的类似光环的通信模式，其中必须仅通信每个图块边缘的数据。cuPyNumeric 是现有唯一支持以这种方式对分布式数组进行叠加和突变的分布式 NumPy 实现。
借助 cuPyNumeric，此 Python 程序可以高效地弱扩展到大量 GPU。图 1 显示了该程序在 NVIDIA Eos 超级计算机上扩展到 1,024 个 GPU 的弱扩展图。NVIDIA Eos 是 2022 年宣布推出的用于推进 AI 研究和开发的超级计算机，由 576 个 NVIDIA DGX H100 节点（共 4,608 个 NVIDIA H100 Tensor Core GPU）组成，通过 400-Gbps 的 NVIDIA Quantum-2 InfiniBand 连接。
弱扩展实验从单个 GPU 上的固定问题大小开始，然后增加 GPU 数量和问题大小，同时保持每个 GPU 的问题大小不变。系统实现的每个 GPU 的吞吐量随后被绘制出来。
图中的近乎平坦的直线表明代码保持在单个 GPU 上实现的相同吞吐量，这意味着当给定更多 GPU 时，cuPyNumeric 可以在相同的时间内解决更大的问题。这是模板计算的理想弱扩展图，显示了 近邻通信模式 ，每个处理器的通信量恒定。
我们还测量了单个 GPU 基准性能，以评估 cuPyNumeric 为分布式执行引入的运行时开销。如图 1 所示，cuPyNumeric 的单个 GPU 性能与基准性能相比，几乎没有分布式执行的开销。
cuPyNumeric 如何并行执行 NumPy 运算
根据设计，NumPy API 围绕具有充足数据并行性的向量运算构建。cuPyNumeric 利用这种固有数据并行性，通过分区数组并使用多个 GPU 对子集并行执行计算。在此期间，当 GPU 对相同的数组子集进行重叠访问时，cuPyNumeric 执行必要的通信。
在本节中，我们将使用 stencil 示例展示 cuPyNumeric 如何并行执行 NumPy 运算。
从计算平均值的语句中提取表达式 center + north。为了在多个 GPU 上并行执行此运算，cuPyNumeric 通过以下方式在 GPU 上划分 center 和 north 数组：如果一个 GPU 的 center 图块包含 center[i,j]，则同一 GPU 的 north 对应图块必须具有 north[i,j]。
对数组进行分区后，cuPyNumeric 会在 GPU 上启动任务。每个 GPU 上的任务在分配给该 GPU 的 center 和 north 的图块上执行元素级添加。图 2 显示了表达式的并行化示例，其中包含四个任务，使用 center 和 north 的 2 x 2 个分区。
即使图 2 绘制得像 center 和 north 是不同的数组，但 center 和 north 实际上是同一数组 grid 的重叠切片，因此将它们映射到单独的物理分配效率很低。
相反，cuPyNumeric 执行 合并优化 ，尝试将同一数组的不同切片使用的数据分组到单个更大的分配中。图 3 显示了图 2 中示例的优化结果。
如果在上一次迭代中更新了 center 数组后再次执行相同的加法，会出现什么情况？
由于 center 和 north 是同一数组的切片，因此在当前迭代中进行加法之前，必须将上一次迭代中对 center 数组所做的更改传播到 north。由于分配合并，每个 GPU 上 center 和 north 之间的交集中的单元已经是最新的。
对于在多个 GPU 上复制的内容（图 4 中的单元 A、B、C 和 D），cuPyNumeric 会自动将数据从 center 复制到 north，以保证额外任务的数据访问一致性。
图 4 显示了图 3 所示示例所需的数据传输。
最后，将范围扩展到平均计算中五个数组的聚合。该代码执行四个元素级添加，每个元素均对应单独的 NumPy API 调用。在完成添加之前，应将上一次迭代中对 center 数组所做的更改分别传播到 north、east、west 和 south 数组。
cuPyNumeric 异步执行此代码，以便通过有用的独立计算来隐藏与通信。为寻找重叠机会，cuPyNumeric 构建了一个 任务图 ，这是一个任务的有向无环图（DAG），其边缘表示任务依赖项。在图形中彼此未连接的节点可以潜在地同时执行。
图 5 显示了内部模板循环的任务图，其中每个加法运算均可用于隐藏数据传输至后续添加所用数组的延迟。为简洁起见，该图形的每个多 GPU 运算只有一个节点。
真实示例中的开发者工作效率
虽然模板示例是一个小型简单的程序，但我们已经能够将大型科学应用程序移植到 cuPyNumeric。
其中一个示例是 TorchSWE ，这是 Pi-Yueh Chuang 和 Lorena Barba 博士开发的 GPU 加速浅水方程求解器。TorchSWE 可求解垂直平均的 Navier-Stokes 方程，并可模拟河流、水道和沿海地区的自由地表水流，以及洪水淹没建模。在给定地形下，TorchSWE 可以预测洪水易发地区和洪水淹没的高度，使其成为风险映射的宝贵工具。
高分辨率数值模拟——例如对需要数亿个数据点的真实地形的模拟——需要跨多个 GPU 进行分布式计算。为此，原始 TorchSWE 实现使用 Mpi4Py 手动划分问题并管理 GPU 间数据通信和同步。
cuPyNumeric 支持 TorchSWE 的分布式实现，从而避免 MPI 实现的复杂性。通过移除所有域分解逻辑将 TorchSWE 移植到 cuPyNumeric 后，无需进一步修改代码，即可跨多个 GPU 和节点轻松扩展。这种可扩展性使用 32 个 GPU 实现了超过 1.2 亿个数据点的高保真模拟，使研究人员无需专门的分布式计算专业知识即可解决洪水淹没建模中的关键科学问题。
总体而言，cuPyNumeric 的实现消除了 20% 的代码，但更重要的是，通过消除了理解和实施分布式执行所需的复杂逻辑的需求，简化了领域科学家的开发和维护。
有关 TorchSWE 示例中 cuPyNumeric 库的生产力方面的更多信息，请参阅 cuPyNumeric 文档中的 TorchSWE 案例研究 。
开始使用
下载 cuPyNumeric 的最简单方法是使用 conda：
$ conda install -c conda-forge -c legate cupynumeric
有关更多信息，请参阅以下资源：
有关 cuPyNumeric 的更多示例，请参阅文档中的 Examples 。有关 cuPyNumeric 的自学教程，请参阅 第 X 章：使用 cuPyNumeric 进行 Distributed Computing 。
尽管 cuPyNumeric 允许对 NumPy 程序进行零代码更改扩展，但并非每个 NumPy 程序都能有效扩展。如需详细了解 cuPyNumeric 的可扩展性能以及可能对性能产生负面影响的常见反模式，请参阅 NVIDIA cuPyNumeric 文档中的 最佳实践 。
总结
本文介绍了 NVIDIA cuPyNumeric，它是 NumPy 的加速和分布式替代品。借助 cuPyNumeric，跨一系列系统（从笔记本电脑到超级计算机）扩展 NumPy 程序就像更改导入语句一样简单。我们还展示了 cuPyNumeric 如何使用模板示例并行执行 NumPy 运算，以实现多 GPU 执行。
如果您有任何意见或疑问，请通过 /nv-legate/discussion GitHub repo 页面或 电子邮件 联系团队。
