OSDI 2022 | 速來(lái)圍觀！微軟亞洲研究院計(jì)算機(jī)系統(tǒng)領(lǐng)域最新論文！

編者按：OSDI（Operating Systems Design and Implementation）是計(jì)算機(jī)系統(tǒng)領(lǐng)域最頂級(jí)的學(xué)術(shù)會(huì)議之一，匯集了全球計(jì)算機(jī)科學(xué)家們對(duì)于計(jì)算機(jī)系統(tǒng)的前瞻性思考。第16屆 OSDI 于2022年7月11日至13日召開(kāi)，本次會(huì)議共有253篇論文投稿，接收49篇，接收率為19.4%。本文中，我們將分享微軟亞洲研究院被 OSDI 2022收錄的3篇論文，希望可以幫助大家了解計(jì)算機(jī)系統(tǒng)領(lǐng)域的前沿趨勢(shì)。歡迎感興趣的讀者閱讀論文原文。

01

Roller: 快速高效的深度學(xué)習(xí)算子編譯器

論文鏈接：https://www.usenix.org/conference/osdi22/presentation/zhu代碼地址：https://github.com/microsoft/nnfusion/
隨著模型規(guī)模的不斷增長(zhǎng)，深度學(xué)習(xí)對(duì)算力的需求也與日俱增。當(dāng)前，深度學(xué)習(xí)硬件加速器（如 GPU、TPU 等）主要依賴算子加速庫(kù)來(lái)支持深度學(xué)習(xí)應(yīng)用。然而，隨著不斷出現(xiàn)的新模型和新硬件類型，業(yè)界對(duì)快速、高效地開(kāi)發(fā)新算子有了更高的要求。算子編譯器（Tensor Compiler）作為一種新途徑，提供了將算子自動(dòng)編譯成對(duì)應(yīng)加速器內(nèi)核代碼的能力。
目前，深度學(xué)習(xí)模型的算子通常被實(shí)現(xiàn)成為多重循環(huán)嵌套的計(jì)算。為了實(shí)現(xiàn)對(duì)這種算子程序性能的進(jìn)一步提升，研究員們通常需要對(duì)已實(shí)現(xiàn)的多重循環(huán)計(jì)算進(jìn)行循環(huán)展開(kāi)、合并、分塊、緩存使用、改變并行度等調(diào)整。這實(shí)質(zhì)上是一個(gè)組合優(yōu)化問(wèn)題，搜索空間巨大，所以主流的算子編譯器往往不得不借助機(jī)器學(xué)習(xí)的方法在這個(gè)巨大的搜索空間中找出較優(yōu)的程序?qū)崿F(xiàn)。而這種搜索過(guò)程往往需要數(shù)千甚至上萬(wàn)步的探索，導(dǎo)致編譯一個(gè)模型經(jīng)常需要耗費(fèi)數(shù)小時(shí)甚至數(shù)天的時(shí)間。這種問(wèn)題在新型的硬件和計(jì)算量巨大的算子上特別嚴(yán)重。因此，目前主流的基于機(jī)器學(xué)習(xí)的算子程序優(yōu)化方法極大地影響了深度學(xué)習(xí)在實(shí)際場(chǎng)景中的部署效率。隨著近年來(lái)深度學(xué)習(xí)模型規(guī)模的快速增長(zhǎng)，預(yù)計(jì)這種情況將進(jìn)一步惡化。
微軟亞洲研究院聯(lián)合多倫多大學(xué)等多所高校在對(duì)深度學(xué)習(xí)計(jì)算中用到的大量算子進(jìn)行編譯、優(yōu)化和性能分析后，發(fā)現(xiàn)了一個(gè)有意思的現(xiàn)象：盡管每個(gè)算子的優(yōu)化選項(xiàng)成千上萬(wàn)，但能達(dá)到最大性能的程序，其配置往往與硬件參數(shù)匹配，從而能夠更充分地利用硬件資源?；谶@樣的觀察，研究員們提出了一個(gè)新的算子編譯器 Roller。為了更好地匹配硬件參數(shù)，Roller 將算子的計(jì)算過(guò)程建模成基于數(shù)據(jù)塊（tile）的“數(shù)據(jù)處理流水線”，即將不同大小的數(shù)據(jù)塊從多級(jí)內(nèi)存結(jié)構(gòu)中搬運(yùn)到計(jì)算核心處理并逐級(jí)寫回。為了使整個(gè)流水線的吞吐最大化，Roller 要求每一級(jí)中數(shù)據(jù)塊的設(shè)置都必須充分匹配硬件的參數(shù)設(shè)置，比如內(nèi)存的訪問(wèn)寬度、計(jì)算核的并行度等等。這樣的約束不僅保證了計(jì)算程序在整個(gè)流水線中的每一級(jí)都充分發(fā)揮了硬件的最大性能，同時(shí)也簡(jiǎn)化了對(duì)算子程序的性能建模，大大減小了搜索空間，從而避免了在實(shí)際硬件加速器上耗時(shí)地進(jìn)行大量嘗試和性能測(cè)量。因此，這樣一種“白盒”的內(nèi)核程序構(gòu)建方法從本質(zhì)上可以大大縮減算子編譯時(shí)間。
圖1：Roller 編譯器的架構(gòu)
更重要的是，Roller 的整個(gè)設(shè)計(jì)是構(gòu)建在一套通用硬件抽象上的，它可以很好地適配到多種主流加速器設(shè)備上。通過(guò)在 NVIDIA GPU、AMD GPU 和 Graphcore IPU 上進(jìn)行實(shí)驗(yàn)評(píng)測(cè)，研究員們發(fā)現(xiàn)相比于當(dāng)前算子編譯器（如 TVM 和 Ansor）和算子庫(kù)（如 cuBLAS、cuDNN 等），Roller 的編譯方法可以生成性能相近，甚至更好的內(nèi)核代碼。例如，研究員們?cè)趯?duì)不同類型的模型中常用到的一百多個(gè)算子進(jìn)行編譯后，Roller 產(chǎn)生的內(nèi)核代碼中分別有59.7%和73.1%的代碼要優(yōu)于 NVIDIA 和 AMD 的算子庫(kù)，有54.6%和58.8%的代碼優(yōu)于 TVM 或 Ansor 在這兩種 GPU 上生成程序的性能。值得注意的是，Roller 的算子編譯時(shí)間也從小時(shí)級(jí)別降低到了秒級(jí)別，編譯時(shí)間減少了三個(gè)數(shù)量級(jí)！
研究員們認(rèn)為 Roller 這樣的系統(tǒng)可以給硬件加速器提供更高效快速的軟件生態(tài)支持。對(duì)于缺少成熟算子庫(kù)的新型硬件廠商，Roller 還提供了一個(gè)可以快速生成算子庫(kù)的機(jī)會(huì)，從而獲得以較快速度趕上領(lǐng)先硬件廠商的機(jī)會(huì)。Roller 的“白盒”編譯方法也給針對(duì)特定硬件快速進(jìn)行深度學(xué)習(xí)模型優(yōu)化和部署帶來(lái)了新的機(jī)遇。微軟亞洲研究院正在基于 Roller 進(jìn)一步完善深度學(xué)習(xí)模型編譯棧，并展開(kāi)更多關(guān)于模型優(yōu)化、場(chǎng)景部署和新型硬件支持的研究。

02

RESIN: 一種處理云服務(wù)中內(nèi)存泄漏的整體方案

論文鏈接：https://www.usenix.org/conference/osdi22/presentation/lou-chang-resin
云計(jì)算的基礎(chǔ)設(shè)施包含許多復(fù)雜的軟件組件，因此會(huì)遇到內(nèi)存泄漏的問(wèn)題。在云服務(wù)中，一旦某個(gè)進(jìn)程發(fā)生內(nèi)存泄露，就會(huì)使得其部署的大量機(jī)器內(nèi)存不足以及無(wú)辜進(jìn)程被殺，最終導(dǎo)致大量用戶的虛擬機(jī)性能下降、崩潰及重啟，嚴(yán)重影響用戶的體驗(yàn)并可能給用戶造成重大經(jīng)濟(jì)損失。
然而，在云系統(tǒng)中，內(nèi)存泄漏問(wèn)題很難處理，現(xiàn)有解決方案或是準(zhǔn)確性不高，或是會(huì)產(chǎn)生很高的開(kāi)銷。其原因，一方面是由于許多內(nèi)存泄漏在極少數(shù)情況下才會(huì)被緩慢觸發(fā)，因此在開(kāi)發(fā)部署過(guò)程中很容易躲過(guò)測(cè)試和檢測(cè)；另一方面，在檢測(cè)到內(nèi)存泄漏后，相關(guān)問(wèn)題很難離線復(fù)現(xiàn)，這會(huì)使得開(kāi)發(fā)人員很難找到泄漏的根本原因。為了解決上述問(wèn)題，微軟亞洲研究院與微軟 Azure 團(tuán)隊(duì)以及霍普金斯大學(xué)的科研人員，一起提出了一種處理云服務(wù)中內(nèi)存泄漏的整體方案 RESIN。
RESIN 采用集中式方法，無(wú)需訪問(wèn)組件的源代碼，也不需要大量的檢測(cè)或重新編譯。RESIN 對(duì)每個(gè)主機(jī)都使用了一個(gè)監(jiān)控代理（Monitor Agent），該代理通過(guò)利用底層操作系統(tǒng)功能來(lái)收集內(nèi)存使用數(shù)據(jù)，因此它支持包括內(nèi)核（Kernel）在內(nèi)的所有組件。收集好的數(shù)據(jù)隨后會(huì)被加載到遠(yuǎn)程服務(wù)用于進(jìn)一步的數(shù)據(jù)分析，從而最大限度地減少主機(jī)的開(kāi)銷。通過(guò)聚合來(lái)自不同主機(jī)的數(shù)據(jù)，RESIN 可以運(yùn)行更復(fù)雜的分析來(lái)捕獲復(fù)雜的泄漏。對(duì)于內(nèi)存泄漏問(wèn)題，RESIN 進(jìn)行了多級(jí)分解和處理。它會(huì)首先執(zhí)行輕量級(jí)泄漏檢測(cè)，并在需要確認(rèn)和診斷時(shí)觸發(fā)更深入的動(dòng)態(tài)檢查。這種分治法（Divide-and-Conquer）使 RESIN 能夠同時(shí)實(shí)現(xiàn)低開(kāi)銷、高精度和可擴(kuò)展性。
圖2：RESIN 系統(tǒng)的工作流程
目前，RESIN 已經(jīng)在微軟 Azure 中運(yùn)行3年，具有高準(zhǔn)確性與低開(kāi)銷等優(yōu)勢(shì)。由于內(nèi)存不足導(dǎo)致的虛擬機(jī)重啟有效地減少了41倍。

03

SparTA: 基于張量稀疏屬性的深度學(xué)習(xí)模型稀疏化編譯框架

論文鏈接：https://www.usenix.org/conference/osdi22/presentation/zheng-ningxin項(xiàng)目代碼地址：https://github.com/microsoft/SparTA.git論文結(jié)果復(fù)現(xiàn)代碼地址：https://github.com/microsoft/nni/tree/sparta_artifact/sparta
隨著深度學(xué)習(xí)的快速發(fā)展，深度學(xué)習(xí)模型的規(guī)模呈現(xiàn)出指數(shù)級(jí)增長(zhǎng)的態(tài)勢(shì)。單個(gè)模型具有高達(dá)萬(wàn)億級(jí)別的參數(shù)量，這遠(yuǎn)遠(yuǎn)超出了硬件加速器運(yùn)算能力的增長(zhǎng)速度。同時(shí)，端側(cè)設(shè)備由于其運(yùn)算能力和功耗的限制，對(duì)深度學(xué)習(xí)模型大小和推理延遲有近乎嚴(yán)苛的要求。因此，探索深度學(xué)習(xí)模型中的稀疏性以及對(duì)稀疏模型做有效的加速成為深度學(xué)習(xí)模型發(fā)展和落地的關(guān)鍵因素。然而，目前稀疏模型在系統(tǒng)支持方面還存在諸多不足，阻礙了模型在稀疏化上的探索。
這些不足具體表現(xiàn)為（i）由于針對(duì)特殊的模型稀疏模式構(gòu)建高效的代碼非常困難，通常需要系統(tǒng)程序員的深度參與，深度學(xué)習(xí)研究員在模型稀疏化的研究過(guò)程中通常使用代理指標(biāo)（Proxy metrics，如 FLOPs、位寬）來(lái)估算加速效果?？墒谴碇笜?biāo)并不能準(zhǔn)確反映模型稀疏化帶來(lái)的加速效果，其預(yù)測(cè)的效果有時(shí)甚至和實(shí)際情況相去甚遠(yuǎn)；（ii）當(dāng)前稀疏優(yōu)化的大多工作都只集中在單個(gè)算子，忽視了一個(gè)稀疏化算子在整個(gè)深度學(xué)習(xí)模型中可能產(chǎn)生的連帶影響，可將稀疏化傳導(dǎo)到模型中的其他算子；（iii）當(dāng)前針對(duì)某個(gè)具體模型的稀疏化優(yōu)化方案很難被重用，難以和其它相關(guān)技術(shù)組合起來(lái)遷移到其他深度學(xué)習(xí)模型上去，例如，在被剪枝的算子和被量化的算子上做的優(yōu)化技術(shù)很難直接遷移到一個(gè)既被剪枝又被量化的算子上。
為了解決模型稀疏化過(guò)程中遇到的上述挑戰(zhàn)，微軟亞洲研究院提出了 SparTA——一個(gè)專門優(yōu)化稀疏深度學(xué)習(xí)模型的端到端編譯框架。SparTA 將深度學(xué)習(xí)模型中張量的稀疏屬性（由剪枝和量化產(chǎn)生）作為整個(gè)編譯框架中的核心抽象 TeSA (Tensor with Sparsity Attribute) ，并圍繞 TeSA 構(gòu)建出面向稀疏模型的全套編譯優(yōu)化技術(shù)。
圖3：SparTA 的系統(tǒng)架構(gòu)
在使用 SparTA 時(shí)，用戶先使用 TeSA 來(lái)標(biāo)注深度學(xué)習(xí)模型中某些張量的稀疏樣式（Sparsity pattern），然后通過(guò) SparTA 提出的三個(gè)核心技術(shù)來(lái)對(duì)模型進(jìn)行端到端地優(yōu)化：一是張量的稀疏屬性在整個(gè)數(shù)據(jù)流圖中的傳播，由 Tensor Algebra 和 Tensor Scrambling 技術(shù)自動(dòng)完成；二是對(duì)稀疏算子進(jìn)行優(yōu)化變換（Transformation），變換為更易于加速、計(jì)算效率更高的稀疏樣式，這種變換使得不同的優(yōu)化技術(shù)可以被有機(jī)地結(jié)合起來(lái)；三是做針對(duì)稀疏張量的代碼特化（Code Specialization），將稀疏樣式硬編碼到算子的代碼中，刪除死代碼（Dead Code），并且針對(duì)給定加速器特化使用其專有硬件（如 sparse tensor core）。
通過(guò)全面的測(cè)試，SparTA 展示出了相比于已有工具高達(dá)平均8.4倍的加速效果。SparTA 不但可以對(duì)稠密模型稀疏化以后進(jìn)行優(yōu)化，而且還能用來(lái)加速一開(kāi)始就采用特定稀疏樣式設(shè)計(jì)的大型預(yù)訓(xùn)練模型。在這方面，SparTA 已經(jīng)被用來(lái)優(yōu)化微軟亞洲研究院研發(fā)的 NUWA 預(yù)訓(xùn)練視頻生成模型。SparTA 對(duì) NUWA 模型中提出的 3DNA 稀疏化注意力機(jī)制達(dá)到了2倍以上的加速。微軟亞洲研究院目前正在對(duì) SparTA 進(jìn)行代碼重整和優(yōu)化，提高易用性，并將盡快開(kāi)源，以促進(jìn)深度學(xué)習(xí)模型稀疏化的研究和實(shí)用化。

*博客內(nèi)容為網(wǎng)友個(gè)人發(fā)布，僅代表博主個(gè)人觀點(diǎn)，如有侵權(quán)請(qǐng)聯(lián)系工作人員刪除。