面向大規模矩陣乘法的編碼計算性能研究

王 艷，王希齡，賴宏達，李念爽

（華東交通大學軟件學院，江西 南昌 330013）

互聯網的飛速發展和大數據處理框架的進步推動了大規模的機器學習部署，并促進了分布式機器學習算法在更多領域中的應用[1-5]。 為了更高效地利用大數據訓練更準確的大模型，在機器學習中引入大量矩陣乘法， 一旦采用了分布式機器學習模式，就意味著大量矩陣乘法的分布式計算。 然而，分布式集群中存在某些計算節點由于各種因素的影響（如節點失效、系統故障、通信瓶頸等），計算速度會以某種隨機的方式變慢，從而使分布式機器學習算法執行時間增加，成為分布式計算系統的主要瓶頸[6]，這種節點被稱作掉隊節點。 目前減緩掉隊節點影響的方法主要是增加某種形式的“計算冗余”，例如，采用副本的方式在多個節點上執行相同的計算任務[7]。 然而，最近的研究結果表明，編碼計算可以更加有效地減輕掉隊節點的影響[8-12]，同時相較于副本方案成本更低。

Tandon 等[13]研究了分布式系統中梯度計算的最優編碼設計， 提出了一種新的編碼計算方案，用于計算函數的和，展示了對梯度進行編碼可以提供同步梯度下降法對失效節點和掉隊節點的容忍。 他們根據工作節點運行速度變慢的程度將掉隊節點分成完全掉隊節點（Full Stragglers）和部分掉隊節點（Partial Stragglers）兩種，針對這兩種情況，提出了一種編碼方案來實現機器學習集群對掉隊節點的魯棒性。 Ye 等[14]提出了通信計算高效梯度編碼方案，它從計算負載、掉隊節點容忍和通信成本3 個方面描述了梯度計算的基本權衡。 Li 等[15]提出了一種MapReduce 的編碼框架，稱為“Coded MapReduce”，它在r（r∈N）個精心選擇的節點上分配每個任務的映射計算，以使網絡內的節點減少r 倍的通信負載。Li 等[16]將文獻[15]中提出的編碼思想應用于Tera-Sort，提出了一種新的分布式排序算法“Coded Tera-Sort”， 大大提高了Hadoop MapReduce 中TeraSort基準的執行時間。 Reisizadeh 等[17]研究由各種不同性能的計算機組成的通用異構分布式計算集群，提出了一個編碼框架， 通過交換冗余來減少計算延遲， 從而加速存在掉隊節點的異構集群的分布式計算。Ferdinandn 等[18]提出了一種用于近似矩陣乘法的任意時刻編碼方案， 通過近似計算的形式來加速分布式計算。 Dutta 等[19]考慮了存在掉隊節點的情況下， 使用并行處理器計算兩個長向量的卷積問題， 提出了存在截止時間的卷積問題的編碼計算方案。 Park 等[20]提出了一個由工作節點組成的層次計算結構。

上述研究都是通過提出編碼方案來增加對掉隊節點的魯棒性，并未對各自編碼方案的性能開銷展開研究，本文重點考察了面向分布式機器學習的各編碼計算方案的任務完成時間和機器計算總時間兩類開銷。 本文定義了面向分布式機器學習的編碼計算的兩個性能指標：一是某個使用該編碼方案的計算任務完成時間，即完成整個計算任務所需要的時間；二是整個計算任務的機器計算總時間，即整個分布式計算系統中所有工作節點對該計算任務進行計算的時間的總和。 通過計算由n 個工作節點組成的分布式系統中第i 個完成計算任務的節點的密度函數，給出工作節點計算任務完成時間符合均勻分布場景下計算任務的完成時間和機器計算總時間的表達式；對比分析了這一場景下3 種應用于矩陣乘法的編碼方案的完成時間和機器計算總時間，并通過實驗對比了指數分布場景下不同情況對任務完成時間與計算節點總計算開銷影響， 提供了計算方案選擇的依據。

1 分布式計算開銷示例

在機器學習算法中，我們經常要執行矩陣乘法的操作。 我們考慮這樣一個矩陣乘法問題：通過輸入矩陣A 和向量X 計算出矩陣B=AX。由于工作節點的內存大小有限，矩陣A 按列被分成3 個大小相等的子矩陣，即A=[A1A2A3]，計算工作分別在圖1和圖2 所示的兩個分布式計算系統中進行的。

圖1 所示的分布式計算系統由1 個主節點和3個工作節點組成，每個工作節點分別預先存儲子矩陣A1，A2，A3中的一個，且各個節點存儲的矩陣互不相同。 主節點將向量X 發送給每個工作節點，工作節點Wi（i=1，2，3）將本地存儲的矩陣與向量X 進行乘法運算。 主節點接收到3 個工作節點的計算結果， 可計算出B。 假設圖1 的計算系統中節點W1，W2和W3的計算任務完成時間（T） 分別為2，6 ms和10 ms，那么第一個計算系統T＝10 ms，機器計算總時間C＝2+6+10=18 ms。

圖1 未編碼方案示例Fig.1 Illustration of uncoded scheme

圖2 所示的分布式計算系統由1 個主節點和5個工作節點組成，使用（5，3）MDS 編碼將矩陣A 分成 的3 個 子 矩 陣 編 碼 為5 個 編 碼 矩 陣A1，A2，A3，A1+A2+A3，A1+2A2+3A3，并分別預先存儲在5 個工作節點W1，W2，W3，W4，W5中。主節點將向量X 發送給每個工作節點，工作節點Wi（i=1，2，3，4，5）將本地存儲的矩陣與向量X 進行乘法運算，并在完成計算任務后，將結果發送回主節點。一旦主節點接收到5個計算結果中的任意3 個，其他工作節點就會停止工作，主節點就可以計算出AX。 例如，當主節點接收到A1X，A3X 和（A1+A2+A3）X 時，它可以通過計算（A1+A2+A3）X-A3X－A1X 得到A2X，然后計算出AX。這一編碼方案，相較于未編碼方案必須接收到全部計算結果才能完成最終計算，能容忍2 個節點計算速度較慢或完全失效。

圖2 編碼方案示例Fig.2 Illustration of coding scheme

假設圖2 中5 個計算節點的完成時間分別為2，4，6，8，10 ms， 主節點在接收到任意3 個節點的計算結果后就立刻通知剩余的其他2 個節點停止計算（根據MDS 編碼性質，任意3 個節點的計算已足夠得到矩陣B），所以任務完成時間T＝6 ms，機器計算總時間C＝2+4+6+6+6=24 ms。

從上述結果可以看出， 相較于未編碼方案，雖然編碼方案的任務完成時間減少了，但是機器計算總時間比未編碼方案更高。

2 模型建立

在一個由1 個主節點和n 個工作節點組成的分布式計算系統中， 由于每個工作節點內存有限，先將總的計算任務分成若干個大小相等的子任務，然后使用編碼技術將這些子任務編碼成n 個大小相同的編碼計算任務，發送給每個工作節點進行計算，主節點在接收到任意k 個工作節點的計算結果后，即可計算出最終結果。

我們通過以下兩個性能指標評估一個編碼方案的性能：

任務完成時間（T）：從主節點將子任務分配給每個工作節點開始，直到主節點接收到足夠多的工作節點提交的計算結果，能夠完成整個計算任務所需要的時間；

機器計算總時間（C）：整個分布式計算系統中所有工作節點對該計算任務進行計算的時間的總和。

假設這n 個工作節點的計算任務完成時間分別為t1，t2， …，tn， 記Xi為這些任務完成時間中第i個最小的完成時間

由于接收到k 個工作節點的計算結果后，主節點即可計算出最終結果，任務完成時間T 為主節點接收到第k 個計算結果的時間，即Xk。 此時分布式計算系統的計算任務就已完成，主節點會向所有工作節點發送一個停止工作的信號，未完成計算任務的（n-k）個工作節點也在Xk時刻停止計算，機器計算總時間C 即為式（2）所示。

由于工作節點在進行計算任務時可能會由于各種原因導致計算速度相較于正常節點更慢或是更快，使得每個工作節點完成計算任務的時間可能會不同，每個工作節點計算任務的完成時間ti（i∈{1，2，…，n}）是一個獨立同分布的連續隨機變量，具有概率分布F 和密度函數f。 為了得到Xi的分布，因Xi小于或等于x 當且僅當這n 個節點的計算時間X1，X2，…，Xn至少有i 個小于或等于x。

3 3 種編碼方案性能的對比分析

對于給定的任何一種編碼策略，本文將完成一個計算任務所需等待的最小工作節點數定義為該編碼策略的恢復閾值。 也就是說，如果任何大小不小于恢復閾值的工作節點子集完成了它們的工作，主節點就能夠計算出最終結果。 在本文中，恢復閾值可以理解為：當主節點接收到這n 個工作節點中最快的前i 個工作節點的計算結果后， 即可計算出最終結果。

在計算各編碼方案的任務完成時間T 和機器計算總時間C 之前，必須要說明的是，由于編碼方案的不同，每種編碼方案中工作節點的計算量是不同的，因而計算時間也不相同。 這在模型中可以體現為，當工作節點的計算時間符合均勻分布的場景下，各個編碼方案的區間（a，b）會有不同。 為了方便，本文將進行一次乘法的計算量記為1，進行加法的計算量忽略不計， 如一個1×n 的向量與一個n×1的向量相乘的計算量為n， 而一個n×1 的向量與一個1×n 的向量相乘的計算量為n2。

為了更好的對比各個編碼方案任務完成時間T和機器計算總時間C， 還需要給這些編碼方案構造一個相同的計算場景，目標都是通過輸入矩陣A 和B 計算出H=ATB，如圖3 所示。矩陣A 的維度為Z×X，矩陣B 的維度為Z×Y，計算任務是在擁有1 個主節點和N 個工作節點的分布式系統中進行。 其中，兩個輸入矩陣分別被（任意）劃分為p×m 和p×n 個子矩陣塊，每一個輸入矩陣劃分的子矩陣大小是相同的，且每個工作節點只能在本地存儲2 個子矩陣。 另外，每個編碼方案的p，m，n 的取值是不同的。

圖3 矩陣乘法問題示例Fig.3 Example of matrix multiplication problem

目前最常用于矩陣乘法的編碼方案為： 一維MDS 編碼、乘積編碼、多項式編碼。本文將對這3 種編碼方案進行介紹以及實驗對比其性能。

3.1 一維MDS 編碼方案

一維MDS 編碼方案是Lee 等[21]對文獻[22]的編碼思想進行擴展得出的， 將MDS 碼在一個輸入矩陣中加入冗余數據這種方法稱為一維MDS 碼（1D MDS 碼）。 該編碼方案的思想是:將大矩陣乘法的問題看作n 個小矩陣乘法的問題，即ATB=［ATb1ATb2… ATbn］， 然后對n 個小矩陣中的每一個分別應用MDS 編碼矩陣乘法。假設N=nk，工作節點被分成大小為k 的n 個組，每個組都專門計算一個ATbj。例如對于第一個組，它用于計算ATb1，該方案首先使用（k，m）MDS 編碼對矩陣A 的沿列分成的m 個大小相同的子矩陣進行編碼， 以獲得k 個編碼矩陣，比如a1到ak，然后將aiTb1的計算分配給這組的第i 個工作節點，以此類推。 MDS 編碼計算的計算時間由n 個組的計算時間中的最大值決定， 每個組的計算時間由組中k 個工作節點中第m 個完成計算任務的工作節點決定。

圖4 所 示 為N=3，m=2，n=1，p=1 時 的1D MDS 編碼方案示例， 主節點接收到3 個工作節點中的任意2 個返回的計算結果后， 即可計算出矩陣H。

圖4 3 工作節點的1D MDS 碼示例Fig.4 Illustration of 1D MDS code with 3 work nodes

3.2 乘積碼編碼方案

圖5 9 工作節點的乘積碼示例Fig.5 Illustration of product code with 9 work nodes

3.3 多項式碼編碼方案

多項式碼編碼方案是Yu 等[23]提出的一種編碼計算策略，利用編碼理論的思想來設計工作節點上的中間計算，以實現對掉隊節點的容忍。 該方案首先沿著列將2 個輸入矩陣分別平均分成m 和n 個大小相等的子矩陣

然后給每個工作節點i（i∈{0，1，…，N-1}）分配一個數字xi，同時使得所有的xi都互不相同，在工作節點i 本地存儲以下2 個編碼子矩陣

由于所設計的計算策略具有特殊的代數結構，解碼過程可以看作是一個多項式插值問題(或是一個解碼Reed-Solomon 碼的問題)，可有效求解。這種多項式編碼的主要創新之處和優點在于，通過精心設計編碼子矩陣的代數結構，可以確保工作節點上的任何mn 個中間計算都足以在主節點上恢復矩陣乘法的最終乘積。 從某種意義上來說，這是在中間計算中創建了一個MDS 結構， 而不是像以前的工作那樣僅僅是對矩陣進行編碼。

圖6 所示為N=5，m=2，n=2，p=1 時的多項式碼編碼方案示例，主節點接收到5 個工作節點中的任意4 個返回的計算結果后，即可計算出矩陣H。

圖6 5 工作節點的多項式碼示例Fig.6 Illustration of polynomial code with 5 work nodes

另外，在多項式編碼的基礎上，Yu 等[24]提出了糾纏多項式編碼方案， 實現了pmn+p-1 的恢復閾值；Dutta 等[25]提出了MatDot 編碼方案，實現了2m-1 的恢復閾值。 由于這幾種方案都是通過設計編碼子矩陣的代數結構實現對掉隊節點的容忍，本文僅將多項式編碼方案與其他編碼方案進行對比分析。

4 實驗及結果

為了更加直觀的比較各個編碼方案的任務完成時間T 和機器計算總時間C，我們通過實驗對3種編碼方案的性能指標進行了測試。

本文使用Python 語言實現了3 種編碼方案的編碼計算過程。 當編碼參數設置為a=1，b=9，X=Y=Z=4，N=16 時， 本文分別對未編碼方案和3 種編碼方案的任務完成時間T 和機器計算總時間C 進行了數值上的測量，實驗結果如表1 所示。

表1 未編碼方案和3 種編碼方案的對比Tab.1 Comparison between uncoded scheme and three coding schemes

從表中結果可以看出，編碼方案與未編碼方案相比：編碼方案的任務完成時間普遍要短，這是因為每種編碼方案只需等待達到其恢復閾值的工作節點子集完成計算任務即可得到最終結果，而不必等待所有節點完成計算任務；但是編碼方案的機器計算總時間普遍更長， 這是因為相較于未編碼方案，編碼方案使用了更多的工作節點來完成計算任務。 將3 種編碼方案進行對比：乘積碼的任務完成時間和機器計算總時間相較于1D MDS 編碼方案都要短，這是由于乘積碼的恢復閾值更低；多項式碼在3 種編碼方案中，任務完成時間和機器計算總時間的表現都是最優，這是因為在給定其他參數情況下3 種編碼方案p，m，n 的取值都相同時，多項式碼的恢復閾值遠小于另外兩種編碼方案，但是該方案額外增加了主節點的編解碼，是以增加主節點的計算時間為代價的。

4.1 編碼方案的參數選擇

選定一個合適的編碼方案之后，比如某用戶準備使用多項式編碼計算方案對2 個大小為100×100的矩陣進行乘法計算，若該用戶選擇將兩個輸入矩陣分別劃分為4 個大小相同的子矩陣， 即m=n=4，假設該用戶選擇在20 個工作節點上執行計算任務，且在該用戶選擇的工作節點的工作環境下a=1，b=9，那么此時任務完成時間T=443 252 ms，機器計算總時間C=4 824 405 ms。 若是該用戶選擇將兩個輸入矩陣分別劃分為4 個大小相同的子矩陣，即m=n=5，假設該用戶選擇在29 個工作節點上執行計算任務，且在該用戶選擇的工作節點的工作環境下a=1，b=9 ， 那么此時任務完成時間T=306 667 ms，機器計算總時間C=4 573 333 ms。 從上述例子中可以看出，不同的參數選擇會導致機器計算總時間C，即資源消耗的不同。 那么我們又面臨著這樣一個新的問題：如何選擇編碼方案的參數，使得用于編碼計算任務的資源消耗最少。

將這一問題使用數學語言可以描述為: 在給定工作節點數量N 和任務完成時間T 的限制時，如何選擇m，n，N 和其他參數的值， 使得機器計算總時間C 最小。

要研究選擇m，n，N 和其他參數的值， 使得機器計算總時間C 最小，我們首先要了解機器計算總時間C 與這些參數之間有什么關系，即參數變化時C 的變化趨勢。 另外，由于對工作節點數量N 和任務完成時間T 的限制可以理解為給定了N 和T 的上限，并不代表著所給工作節點數量N 和任務完成時間T 的值就是最終結果，我們還需要對工作節點數量N 和任務完成時間T 進行研究，分析這兩個參數之間、兩個參數與機器計算總時間C 之間以及兩個參數與其他參數之間的關系。

我們首先對前文求出的3 種編碼方案的任務完成時間T 和機器計算總時間C 的表達式進行研究。 從3 種編碼方案性能指標的表達式中，可以很明顯地觀察到，當其他參數保持不變時，增加Z，X，Y 的值，T 和C 都會隨之增加；增加mn 的值，T 和C都會隨之減小。 但是當N 增加時，T 和C 的值不容易判斷，本文又對工作節點數N 變化的情況下3 種編碼方案的T 和C 的變化情況進行了實驗分析。在其他參數分別設置為a=1，b=9，X=Y=Z=100，m=n=5，p=1 時，計算出N 分別取值為30，35，40，45，50，55，60，65，70，80，90，100，120，160，200，240，280，320，360 時各編碼方案的T 和C 的值，如圖7 所示。

圖7 3 種編碼方案在不同N 取值下的T 和C 變化情況Fig.7 Changes of T and C in three coding schemes with different values of N

從圖7 中可以看出， 隨著工作節點數N 的增加，3 種編碼方案的機器計算總時間C 都在不斷增加，而任務完成時間T 都在不斷減小，這說明增加更多的冗余工作節點，能更快的完成計算任務，但是計算任務成本開銷也會相應地增加。 另外，從圖中還可以看出，在C 增加到一定程度后， T 的減小會變得緩慢， 這說明增加的冗余達到一定程度之后， 再繼續增加冗余， 成本開銷依然會繼續增加， 但是任務完成時間卻減少的很慢甚至不再減少，此時收益很低。

當輸入矩陣A 和B 劃分成的子矩陣數量不同，即參數m，n 不同時，我們通過實驗對編碼方案的任務完成時間T 和機器計算總時間C 進行了研究。 以多項式碼編碼方案為例，當a=1，b=9，X=Y=Z=4 時，本文分別給出了m，n 取值不同的情況下T 和C 的變化情況，研究結果如圖8 和圖9 所示。

圖8 m 和n 值不同時多項式編碼框架的T 的變化情況Fig.8 Changes of T for polynomial code scheme with different values of m and n

圖9 m 和n 值不同時多項式編碼框架的C 的變化情況Fig.9 Changes of C for polynomial code scheme with different values of m and n

可以看出，當m 和n 的取值越大時，編碼子矩陣就越復雜，生成編碼子矩陣以及對編碼子矩陣進行編解碼所需要的時間也就越長。 m 和n 的取值越大并不能使進行編碼計算任務花費的全部時間越短。

從圖9 可以看出，當給定N 時，選擇m×n 較大的輸入矩陣劃分方案，可以得到較小的C。這是因為當工作節點的總數量保持不變時，每個工作節點的工作時間減少， 機器計算總時間自然也就隨之減少。 此外，當工作節點數量N 與多項式碼的恢復閾值mn 相等時， 也就是說在分布式計算系統中沒有冗余存在，此時無論N 的取值如何變化，多項式編碼方案的機器計算總時間C 總是相等的，這是由C的計算表達式決定的，這也是多項式編碼方案的機器計算總時間C 如圖7 變化的原因。

4.2 參數的啟發式算法

在對4.1 節實驗結果進行分析的基礎上， 本文給出了在給定T 和N 的上限時， 如何選取m，n，T的值來最小化機器計算總時間C 的啟發式算法，算法過程如下所示：

輸入：矩陣維度X，Y，Z，計算節點完成最快時間a，最慢時間b，給定T 和N 的上限；

輸出：計算總時間C 的預測值，及最小化C 的參數m，n 的取值，任務完成時間T 和工作節點數量N；

1） 根據輸入矩陣大小求出所有可用的矩陣劃分方案m 和n 的可能值；

2） 根據用戶所給的任務完成時間T 的限制選擇合適的m 和n 的取值；

3） 根據工作節點數量N 的上限判斷在給定T和上限內是否能完成計算工作；

4） 若能則給出使得機器計算總時間C 最小的各項參數的取值。

在算法中輸入給定的任務完成時間T 和工作節點數量N 的上限以及需要進行乘法計算的矩陣大小， 然后算法會輸出可以使得機器計算總時間C最小化的參數m，n 的取值，任務完成時間T 和工作節點數量N 的預測值以及機器計算總時間C 的預測值。

當算法輸入為a=1，b=9，X=Y=Z=100，N=46，T=200 000 ms 時，輸出為m=2，n=20 或m=4，n=10 或m=10，n=4 或m=20，n=2，T=195 212.8 ms，C=5 660 638.3 ms。

在mn 不同取值下C 的取值如表2 所示， 可以看出算法結果的確可以輸出使得機器計算總時間C最小化的參數m，n 的取值。

表2 在mn 不同取值下C 的值Tab.2 C with different values of m×n

5 結論

1） 本文對基于分布式機器學習的編碼方案的性能開銷進行了研究和分析，計算出了n 個工作節點中第i 個完成計算任務的節點的密度函數，并給出了工作節點計算任務完成時間符合均勻分布場景下計算任務的任務完成時間C 和機器計算總時間T 的表達式。

2） 分析對比了3 個應用于矩陣乘法的編碼方案與未編碼方案的任務完成時間和機器計算總時間，提供了計算方案選擇的依據。 并且在此基礎上，通過Python 實驗了指數分布情況下的編碼方案并實驗得出了其任務完成時間C 和機器計算總時間T，還在此基礎上進行實驗，對3 種編碼方案進行比較。

3） 在此基礎上還對參數選擇進行分析， 提出一個啟發式算法為參數的選擇提供參考。