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 npimport cupynumeric as npnum_points = 1000000000# Generates random points in a square of area 4x = 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 1probability_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 npgrid = 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 頁面或 電子郵件 聯系團隊。
?
