新聞中心

借助 NVIDIA cuPyNumeric,將 NumPy 無縫擴展從筆記本電腦到超級計算機 NEWS DETAIL

當前位置:首頁 > 新聞中心
資訊分類 · 新聞中心 發佈時間 · 2025-01-17

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 之後,中心數組的更新必須傳播到 northeastwest 和 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 性能與基準性能相比,幾乎沒有分布式執行的開銷。

The plot demonstrates near-perfect scalability out to 128 GPUs on the NVIDIA Eos supercomputer.
圖 1. Weak scaling plot for a cuPyNumeric 2D stencil computation

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 個分區。

A diagram shows the original grid, which is first split into center and north sections. Each section is then further divided across four GPUs.
圖 2 、 四個 GPU 上 表達式中心+北並行化的示例(PyTorch 上的 Megatron-LM 例子)

即使圖 2 繪制得像 center 和 north 是不同的數組,但 center 和 north 實際上是同一數組 grid 的重疊切片,因此將它們映射到單獨的物理分配效率很低。

相反,cuPyNumeric 執行 合併優化 ,嘗試將同一數組的不同切片使用的資料分組到單個更大的分配中。圖 3 顯示了圖 2 中示例的優化結果。

The diagram shows how the center and north partitions overlap when they are originally partitioned across four GPUs.
圖 3. 前一個示例中的 中心 和 北側合併

如果在上一次迭代中更新了 center 數組後再次執行相同的加法,會出現甚麼情況?

由於 center 和 north 是同一數組的切片,因此在當前迭代中進行加法之前,必須將上一次迭代中對 center 數組所做的更改傳播到 north。由於分配合併,每個 GPU 上 center 和 north 之間的交集中的單元已經是最新的。

對於在多個 GPU 上複製的內容(圖 4 中的單元 ABC 和 D),cuPyNumeric 會自動將資料從 center 複製到 north,以保證額外任務的資料訪問一致性。

圖 4 顯示了圖 3 所示示例所需的資料傳輸。

A diagram shows how data transfer from the center array to the north array is implemented across the four partitions of the original grid.
圖 4 、通過資料傳輸將變化 從中心傳播到北部

最後,將範圍擴展到平均計算中五個數組的聚合。該代碼執行四個元素級添加,每個元素均對應單獨的 NumPy API 調用。在完成添加之前,應將上一次迭代中對 center 數組所做的更改分別傳播到 northeastwest 和 south 數組。

cuPyNumeric 異步執行此代碼,以便通過有用的獨立計算來隱藏與通信。為尋找重疊機會,cuPyNumeric 構建了一個 任務圖 ,這是一個任務的有向無環圖(DAG),其邊緣表示任務依賴項。在圖形中彼此未連接的節點可以潛在地同時執行。

圖 5 顯示了內部模板循環的任務圖,其中每個加法運算均可用於隱藏資料傳輸至後續添加所用數組的延遲。為簡潔起見,該圖形的每個多 GPU 運算只有一個節點。

A diagram includes four boxes corresponding to NumPy API calls, and four more boxes with arrows pointing to the first four, representing data transfers inserted by cuNumeric. The arrows between the boxes indicate dependencies between the various operations.
圖 5. 用於求和中心、北、東、西和南數組的任務圖  綠色 = NumPy API 調用; 灰色 cuPyNumeric 插入的資料傳輸)

真實示例中的開發者工作效率 

雖然模板示例是一個小型簡單的程序,但我們已經能夠將大型科學應用程序移植到 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

有關更多資訊,請參閱以下資源:

  • NVIDIA cuPyNumeric User Guide

    • 安裝

  • /nv-legate/cupynumeric GitHub 庫

有關 cuPyNumeric 的更多示例,請參閱文檔中的 Examples 。有關 cuPyNumeric 的自學教程,請參閱 第 X 章:使用 cuPyNumeric 進行 Distributed Computing 

儘管 cuPyNumeric 允許對 NumPy 程序進行零代碼更改擴展,但並非每個 NumPy 程序都能有效擴展。如需詳細瞭解 cuPyNumeric 的可擴展性能以及可能對性能產生負面影響的常見反模式,請參閱 NVIDIA cuPyNumeric 文檔中的 最佳實踐 

總結 

本文介紹了 NVIDIA cuPyNumeric,它是 NumPy 的加速和分布式替代品。借助 cuPyNumeric,跨一系列系統(從筆記本電腦到超級計算機)擴展 NumPy 程序就像更改導入語句一樣簡單。我們還展示了 cuPyNumeric 如何使用模板示例並行執行 NumPy 運算,以實現多 GPU 執行。