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

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

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

庫把手

庫句柄表示參與計(jì)算的設(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類型的錯(cuò)誤代碼。在生產(chǎn)中，您應(yīng)該始終檢查錯(cuò)誤代碼，以便盡早檢測故障或使用問題。為了簡潔起見，我們在本文中省略了這些檢查，但它們包含在相應(yīng)的示例代碼中。

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

張量描述符

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

extent：每個(gè)模式的邏輯大小。

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è)備分配到使用另一個(gè)數(shù)組devices編碼的塊。該陣列是一個(gè)密集的柱狀主張量，其范圍與設(shè)備計(jì)數(shù)類似。

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

最后，設(shè)備上的確切數(shù)據(jù)布局由每種模式的elementStride和blockStride 值決定。它們分別以元素為單位在線性存儲器中確定給定模式下兩個(gè)相鄰元素和相鄰塊的位移（圖 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)生交錯(cuò)塊格式。如果deviceCount為NULL，則所有設(shè)備計(jì)數(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）。這兩個(gè)張量在重合模式下的范圍必須匹配，但它們的其他方面可能不同。一個(gè)可能位于主機(jī)上，另一個(gè)跨設(shè)備，它們可能具有不同的阻塞和步幅。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

收縮手術(shù)

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

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

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

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

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

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

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

cutensorMgCreateContractionFind(handle, &find, CUTENSORMG_ALGO_DEFAULT);

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

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

創(chuàng)建一個(gè)計(jì)劃：

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

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

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

在這個(gè)調(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ò)展為兩個(gè)參數(shù)： GPU 的數(shù)量和比例因子。您可以隨意嘗試其他收縮、塊大小和縮放模式。它是以這樣一種方式編寫的，即在保持 K 不變的情況下，將 M 和 N 放大。它實(shí)現(xiàn)了形狀的幾乎 GEMM 形狀的張量收縮：

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

關(guān)于作者

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

審核編輯：郭婷