如何測量各種工作負(fù)載和GPU配置下收縮操作的性能

張量收縮是機(jī)器學(xué)習(xí)、計算化學(xué)和量子計算中許多重要工作的核心。隨著科學(xué)家和工程師們對不斷增長的問題的研究，基礎(chǔ)數(shù)據(jù)變得越來越大，計算時間也越來越長。

當(dāng)張量收縮不再適合單個 GPU 時，或者如果在單個 GPU 上花費(fèi)的時間太長，自然下一步是將收縮分布到多個 GPU 上。我們一直在用這個新功能擴(kuò)展 cuTENSOR ，并將其作為一個名為 cuTENSORMg （多 GPU ）的新庫發(fā)布。它在塊循環(huán)分布張量上提供單進(jìn)程多 GPU 功能。

cuTENSORMg 的copy和contraction操作大致分為句柄、張量描述符和描述符。在這篇文章中，我們將解釋句柄和張量描述符，以及復(fù)制操作是如何工作的，并演示如何執(zhí)行張量收縮。然后，我們將展示如何測量各種工作負(fù)載和 GPU 配置下收縮操作的性能。

庫把手

庫句柄表示參與計算的設(shè)備集。句柄還包含跨調(diào)用重用的數(shù)據(jù)和資源。通過將設(shè)備列表傳遞給cutensorMgCreate函數(shù)，可以創(chuàng)建庫句柄：

cutensorMgCreate(&handle, numDevices, devices);

cuTENSORMg 中的所有對象都是堆分配的。因此，必須通過匹配的destroy調(diào)用釋放它們。為了簡潔起見，我們在這篇文章中沒有展示這些，但是生產(chǎn)代碼應(yīng)該銷毀它創(chuàng)建的所有對象，以避免泄漏。

cutensorMgDestroy(handle);

所有庫調(diào)用都返回cutensorStatus_t類型的錯誤代碼。在生產(chǎn)中，您應(yīng)該始終檢查錯誤代碼，以便盡早檢測故障或使用問題。為了簡潔起見，我們在本文中省略了這些檢查，但它們包含在相應(yīng)的示例代碼中。

除了錯誤代碼， cuTENSORMg 還提供與 cuTENSOR 類似的日志記錄功能 ?？梢酝ㄟ^適當(dāng)設(shè)置CUTENSORMG_LOG_LEVEL環(huán)境變量來激活這些日志。例如，CUTENSORMG_LOG_LEVEL=1將為您提供有關(guān)返回的錯誤代碼的附加信息。

張量描述符

張量描述符描述了張量在內(nèi)存中的布局以及在設(shè)備中的分布。對于每種模式，有三個核心概念來確定布局：

extent：每個模式的邏輯大小。

blockSize：將extent細(xì)分為大小相等的塊，但最后的剩余塊除外。

deviceCount：確定塊在設(shè)備上的分布方式。

圖 1 顯示了extent和block size如何細(xì)分二維張量。

圖 1 帶有范圍和塊的張量數(shù)據(jù)布局。綠色表示二維張量，藍(lán)色表示塊大小導(dǎo)致的塊。

圖 2 以塊循環(huán)方式在設(shè)備上分布塊張量；不同的顏色代表不同的設(shè)備。

塊以循環(huán)方式分布，這意味著連續(xù)的塊被分配給不同的設(shè)備。圖 2 顯示了塊到設(shè)備的逐塊分布，將設(shè)備分配到使用另一個數(shù)組devices編碼的塊。該陣列是一個密集的柱狀主張量，其范圍與設(shè)備計數(shù)類似。

圖 3 使用元素步距和塊步距的設(shè)備上數(shù)據(jù)布局。

最后，設(shè)備上的確切數(shù)據(jù)布局由每種模式的elementStride和blockStride 值決定。它們分別以元素為單位在線性存儲器中確定給定模式下兩個相鄰元素和相鄰塊的位移（圖 3 ）。

這些屬性都是使用cutensorMgCreateTensorDescriptor調(diào)用設(shè)置的：

cutensorMgCreateTensorDescriptor(handle, &desc, numModes, extent, elementStride, blockSize, blockStride, deviceCount, numDevices, devices, type);

可以將NULL傳遞給elementStride、blockSize、blockStride和deviceCount。

如果elementStride是NULL，則使用通用列主布局假定數(shù)據(jù)布局密集。如果blockSize是NULL，則等于extent。如果blockStride是NULL，則它等于blockSize * elementStride，這將產(chǎn)生交錯塊格式。如果deviceCount為NULL，則所有設(shè)備計數(shù)都設(shè)置為 1 。在這種情況下，張量是分布式的，完全駐留在devices［0］的內(nèi)存中。

通過將CUTENSOR_MG_DEVICE_HOST作為所屬設(shè)備傳遞，可以指定 tensor 位于主機(jī)上的固定、托管或定期分配的內(nèi)存中。

復(fù)制操作

copy操作可以更改數(shù)據(jù)布局，包括將張量重新分配到不同的設(shè)備。其參數(shù)是源和目標(biāo)張量描述符（descSrc和descDst），以及源和目標(biāo)模式列表（modesSrc和modesDst）。這兩個張量在重合模式下的范圍必須匹配，但它們的其他方面可能不同。一個可能位于主機(jī)上，另一個跨設(shè)備，它們可能具有不同的阻塞和步幅。

與 cuTENSORMg 中的所有操作一樣，它分三步進(jìn)行：

cutensorMgCopyDescriptor_t：編碼應(yīng)該執(zhí)行的操作。

cutensorMgCopyPlan_t：編碼操作的執(zhí)行方式。

cutensorMgCopy：根據(jù)計劃執(zhí)行操作。

第一步是創(chuàng)建復(fù)制描述符：

cutensorMgCreateCopyDescriptor(handle, &desc, descDst, modesDst, descSrc, modesSrc);

有了拷貝描述符，您可以查詢所需的設(shè)備端和主機(jī)端工作空間的數(shù)量。deviceWorkspaceSize陣列的元素數(shù)量與手柄中的設(shè)備數(shù)量相同。i-th 元素是句柄中i-th 設(shè)備所需的工作空間量。

cutensorMgCopyGetWorkspace(handle, desc, deviceWorkspaceSize, &hostWorkspaceSize);

確定工作空間大小后，規(guī)劃副本。你可以傳遞一個更大的工作空間大小，呼叫可能會利用更多的工作空間，或者你可以嘗試傳遞一個更小的大小。規(guī)劃可能能夠適應(yīng)這一點，否則可能會產(chǎn)生錯誤。

cutensorMgCreateCopyPlan(handle, &plan, desc, deviceWorkspaceSize, hostWorkspaceSize

最后，計劃完成后，執(zhí)行copy操作。

cutensorMgCopy(handle, plan, ptrDst, ptrSrc, deviceWorkspace, hostWorkspace, streams);

在這個調(diào)用中，ptrDst和ptrSrc是指針數(shù)組。它們包含對應(yīng)的張量描述符中每個設(shè)備的一個指針。在本例中，ptrDst[0]對應(yīng)于作為devices[0]傳遞給cutensorMgCreateTensorDescriptor的設(shè)備。

另一方面，deviceWorkspace和streams也是數(shù)組，其中每個條目對應(yīng)一個設(shè)備。它們是根據(jù)庫句柄中設(shè)備的順序排序的，例如deviceWorkspace[0]和streams[0]對應(yīng)于在devices[0]傳遞給cutensorMgCreate的設(shè)備。工作空間必須至少與傳遞給cutensorMgCreateCopyPlan的工作空間大小相同。

收縮手術(shù)

cuTENSORMg 庫的核心是contraction操作。它目前實現(xiàn)了一個或多個設(shè)備上張量的張量收縮，但將來可能支持主機(jī)上的張量。作為復(fù)習(xí)，收縮是以下形式的操作：

其中、、和是張量，、、和是可以任意排列和交錯的模式列表。

與copy操作一樣，它分三個階段進(jìn)行：

• cutensorMgCreateContractionDescriptor：對問題進(jìn)行編碼。
• cutensorMgCreateContractionPlan：對實現(xiàn)進(jìn)行編碼。
• cutensorMgContraction：使用計劃并執(zhí)行實際收縮。

首先，根據(jù)張量描述符、模式列表和所需的計算類型（例如計算期間可能使用的最低精度數(shù)據(jù)）創(chuàng)建收縮描述符。

cutensorMgCreateContractionDescriptor(handle, &desc, descA, modesA, descB, modesB, descC, modesC, descD, modesD, compute);

由于收縮操作有更多的自由度，您還必須初始化find對象，以便更好地控制給定問題描述符的計劃創(chuàng)建。目前，這個find對象只有一個默認(rèn)設(shè)置：

cutensorMgCreateContractionFind(handle, &find, CUTENSORMG_ALGO_DEFAULT);

然后，您可以按照為copy操作所做的操作來查詢工作空間需求。與該操作相比，您還傳入了find和workspace首選項：

cutensorMgContractionGetWorkspace(handle, desc, find, CUTENSOR_WORKSPACE_RECOMMENDED, deviceWorkspaceSize, &hostWorkspaceSize);

創(chuàng)建一個計劃：

cutensorMgCreateContractionPlan(handle, &plan, desc, find, deviceWorkspaceSize, hostWorkspaceSize);

最后，使用計劃執(zhí)行收縮：

cutensorMgContraction(handle, plan, alpha, ptrA, ptrB, beta, ptrC, ptrD, deviceWorkspace, hostWorkspace, streams);

在這個調(diào)用中， alpha 和 beta 是與張量類型相同的主機(jī)指針，除非張量是半精度或BFloat16精度，在這種情況下是single precision。不同數(shù)組ptrA、ptrB、ptrC和ptrD中指針的順序?qū)?yīng)于它們在描述符devices數(shù)組中的順序。deviceWorkspace和streams數(shù)組中指針的順序與庫句柄的devices數(shù)組中的順序相對應(yīng)。

表演

你可以在CUDA 庫樣本GitHub 回購。我們將其擴(kuò)展為兩個參數(shù)： GPU 的數(shù)量和比例因子。您可以隨意嘗試其他收縮、塊大小和縮放模式。它是以這樣一種方式編寫的，即在保持 K 不變的情況下，將 M 和 N 放大。它實現(xiàn)了形狀的幾乎 GEMM 形狀的張量收縮：

M1和N1按比例放大，這些尺寸中的塊大小保持負(fù)載大致平衡。下圖顯示了在 DGX A100 上測量時的比例關(guān)系。

關(guān)于作者

Markus Hoehnerbach 是 cuTENSOR 和 cuTENSORMg 的高級軟件工程師。他擁有 RWTH 亞琛大學(xué)計算機(jī)科學(xué)博士學(xué)位。他感興趣的領(lǐng)域是結(jié)構(gòu)化和非結(jié)構(gòu)化張量的高性能計算及其在機(jī)器學(xué)習(xí)和計算科學(xué)中的應(yīng)用。

審核編輯：郭婷