新技術高效加速稀疏張量，為大型 AI 模型帶來巨大提升
作者：亞當·西維（Adam Zewe），麻省理工學院
麻省理工學院（MIT）和 NVIDIA 的研究人員開發出了兩種互補的技術，可以顯著提高圖分析和生成式人工智慧等高效能計算應用的速度和效能。這兩種創新方法旨在高效地利用稀疏張量（零值）的特性。
張量是機器學習模型使用的資料結構。這兩種新方法旨在高效地利用張量中的零值。在處理這些張量時，可以跳過零值，節省計算和記憶體。例如，任何與零相乘的數值都是零，所以可以跳過該操作。並且可以壓縮張量（零值不需要儲存），以便更多部分可以儲存在片上記憶體中。然而利用稀疏性仍然存在一些挑戰。在較大的張量中查詢非零值並不容易。現有的方法通常透過強制引入稀疏模式來約束非零值的位置，以簡化搜尋。但這約束了可以高效處理的稀疏張量的種類。另一個挑戰是張量不同區域的非零值數量可能有所不同。這使得難以確保在記憶體中儲存不同區域所需的空間大小。為了確保區域適合，通常會分配比實際所需更多的空間，這會增加片外記憶體流量，需要額外的計算。
解決稀疏性挑戰的兩種方法
麻省理工學院和 NVIDIA 的研究人員提出了兩種解決這些問題的方法。
第一種方法：高效尋找零值
稀疏性可能是由於多種原因引起的。例如，研究人員有時會透過將張量中的某些值替換為零來修剪機器學習模型中的不必要部分，以創造稀疏性。不同模型的稀疏程度（零值的百分比）和零值的位置可能不同。為了更容易尋找模型中剩餘的非零值，研究人員通常會約束非零值的位置，使其遵循特定的模式。然而每個硬體加速器通常只能支援一種特定的稀疏模式，這約束了其靈活性。麻省理工學院開發的高亮（HighLight）硬體加速器可以處理多種稀疏模式，並在執行不包含任何零值的模型時仍然表現出色。他們使用了一種稱為「分級結構稀疏性」的技術，可以高效地表示由多個簡單稀疏模式組成的各種稀疏模式。該方法將張量中的數值劃分為較小的塊，每個塊都有自己的簡單稀疏模式（例如，四個值中有兩個零和兩個非零）。然後，將塊組合成一個層次結構，每個塊集合也具有自己的簡單稀疏模式（例如，一個零塊和三個非零塊在四個塊的層次中）。他們繼續將塊組合成較大的層次，但每一步的模式都是簡單的。這種簡單性使得 HighLight 更容易找到和跳過零值，因此可以充分利用減少多餘計算的機會。平均而言，他們的加速器設計比其他方法節能約六倍。
他們的目標是將分級結構稀疏性應用於更多型別的機器學習模型和模型中的不同型別的張量。
第二種方法：有效“超預訂”以加速工作負載
研究人員還利用稀疏性，更高效地移動和處理計算機晶片上的資料。由於張量的大小通常超過晶片上的記憶體緩衝區大小，因此晶片一次只能抓取和處理一個張量的一個片段，稱為磚塊（tiles）。為了最大程度利用緩衝區的利用率，並減少晶片必須存取片外記憶體的次數，而片外記憶體的存取通常是能源消耗的主要來源以及約束處理速度的原因，研究人員試圖使用最大的可容納的磚塊大小。但在稀疏張量中，很多資料值都是零，因此根據緩衝區容量，甚至可以放入比預期更大的磚塊。不需要儲存零值。但是每個磚塊中的零值數量可能會在張量的不同區域有所不同，因此對於每個磚塊，零值也可能不同。這使得很難確保一個符合緩衝區大小的磚塊大小。因此現有的方法通常保守地假設沒有零值，最終選擇一個較小的磚塊大小，這導致緩衝區中浪費了部分空間。為理解決這個不確保性，研究人員提出了使用“超預訂”來允許他們增加磚塊的大小以及一種容忍磚塊不適合緩衝區的方式。就像航空公司超售航班機票一樣，如果所有乘客都出現，航空公司必須對被撤銷的乘客進行補償。但通常並非所有乘客都會出現。在稀疏張量中，可以選擇磚塊大小，以使大多數磚塊仍然具有足夠的零值以適應緩衝區。但偶爾，某個磚塊的非零值可能超出了可容納量。在這種情況下，這些資料會從緩衝區中被移除。研究人員使硬體僅重新提取被移除的資料而無需再次抓取和處理整個磚塊。他們修改了緩衝區的“尾部”來實現這一點，因此這種技術被命名為“Tailors”。然後，他們還建立了一種磚塊大小的方法，利用超預訂這一概念。這種方法稱為“Swiftiles”，快速估算出理想磚塊大小，以使設定的百分比的磚塊進行超預訂。Swiftiles 減少了硬體需要檢查張量以識別理想磚塊大小的次數，節省了計算量。Tailors 和 Swiftiles 的組合使速度增加了一倍以上，同時只需要現有無法處理超預訂的硬體加速器所需的一半能量。
研究人員希望將超預訂的概念應用於計算機架構的其他方面，並改進評估最佳超預訂水平的流程。
結論
這項研究的成果，由麻省理工學院和 NVIDIA 的研究人員共同完成，使用了兩種互補的技術，顯著提高了稀疏張量的處理速度和能源效率。這些技術為大型 AI 模型帶來了巨大的提升，並擁有良好的靈活性和適應性，同時保持了專門化和高效利用硬體加速器的特點。
這項研究的結果已經在 arXiv 預印本服務上發表，並將在 IEEE/ACM 國際微架構研討會上進行發表。這些突破性的技術將有助於改進高效能計算應用，並為未來的 AI 發展提供更多可能性。
（本文根據麻省理工學院提供的原始新聞稿編寫，內容來源受到版權保護，僅供訊息參考）
