TPU的起源,Jeff Dean綜述後摩爾定律時代的ML硬件與算法

選自arxiv

深度學習和硬件怎樣結合?Jeff Dean 長文介紹了後摩爾定律時代的機器學習研究進展,以及他對未來發展趨勢的預測判斷

過去十年我們見證了機器學習的顯著進步,特別是基於深度學習的神經網絡。機器學習社區也一直在嘗試構建新模型,用於完成具有挑戰性的工作,包括使用強化學習,通過和環境進行交互的方式完成難度較大的任務,如下圍棋、玩電子遊戲等。

機器學習對算力的需求無疑是龐大的,從計算機視覺到自然語言處理,更大的模型和更多的數據往往能夠取得更好的性能。在摩爾定律時代,硬件進步帶來的算力增長尚且能夠滿足機器學習的需求,但當摩爾定律被榨乾後,怎樣讓硬件中的算力資源被機器學習模型充分利用成了下一個需要探討的問題。

TPU的起源,Jeff Dean綜述後摩爾定律時代的ML硬件與算法

文章地址:https://arxiv.org/abs/1911.05289

在這篇文章中,Jeff Dean 首先討論了機器學習的進步對計算設備的影響,特別是在後莫爾定律時代,他詳細介紹了谷歌開發張量運算單元的初衷與效果。與此同時,Jeff Dean 還討論了深度學習如何能夠幫助設計電路與芯片。最終,Jeff Dean 描述了將來機器學習發展的方向——更大尺度的多任務模型,它能夠根據任務要求選擇性地激活,而且可以有更動態的調整過程。

整體文章比較長,機器之心簡要概括了不同章節的主要思想,感興趣的讀者可以在 arXiv 上查閱原文章。
摩爾定律、後摩爾定律和機器學習的計算需求

自 20 世紀 60 年代至 90 年代,基於深度學習和人工神經網絡的諸多關鍵思想和算法已經出現,並且在 80 年代末至 90 年代初,人們開始意識到神經網絡能夠以有趣的方式解決一些問題,並且它們的實質性優勢在於既可以接受原始形式的輸入數據並且在訓練模型執行預測任務的過程中可以自動增強模型的層次表徵,因而激發了 ML 和 AI 社區的研究興趣。

但是,這種方法受限於當時不夠強大的計算機性能。一些研究試圖利用並行算法來擴展訓練神經網絡的計算量,但在大多數情況下,ML 和 AI 社區的關注重點已經不在基於神經網絡的方法了。直到 2005 年以後,在摩爾定律推動計算性能提升的 20 年後,計算機性能才強大到可以訓練大型神經網絡來解決 ImageNet 等現實世界的實際問題,而不再侷限於 MNIST 等小規模問題。其中,GPU 卡上的通用計算範式使得神經網絡開始在解決棘手的實際問題上展示有趣的結果。

TPU的起源,Jeff Dean綜述後摩爾定律時代的ML硬件與算法

摩爾定律提出者戈登·摩爾。

然而令人遺憾的是,在計算性能開始足夠強大到可以解決有趣的現實問題,並且機器學習規模的擴大和適用性的提升也迫切需要更多的計算資源來解決更大的問題,這時計算機行業卻隨著通用 CPU 性能的持續提升出現了整體性的大幅度放緩。下圖 2 展示了近 40 年間的計算性能增長態勢,其中自 1985 年至 2003 年,通用 CPU 性能每 1.5 年提升一倍;自 2003 年至 2010 年,通用 CPU 性能每 2 年提升一倍;而 2010 年以後,通用 CPU 性能預計每 20 年才能提升一倍。

TPU的起源,Jeff Dean綜述後摩爾定律時代的ML硬件與算法

圖 2:摩爾定律和後摩爾定律時代的計算需求增長態勢。

下圖 3 展示了近年來出現的一些重要的機器學習網絡架構,如 AlexNet、GoogleNet、AlphaZero 等,它們對計算的需求呈現越來越大的趨勢。

TPU的起源,Jeff Dean綜述後摩爾定律時代的ML硬件與算法

圖 3:AlexNet、GoogleNet、AlphaZero 等重要的機器學習網絡架構以及它們的計算需求增長態勢。

下圖 4 展示了 2009 年以來,機器學習領域 Arxiv 發表論文數量的增長情況,其中 2018 年的論文數量是 2009 年的 32 倍(論文數量每 2 年提升一倍以上)。現在,每天仍有 100 多篇論文在 Arxiv 機器學習子板塊發表,並且這一增長態勢沒有出現放緩的跡象。

TPU的起源,Jeff Dean綜述後摩爾定律時代的ML硬件與算法

圖 4:自 2009 年以來,機器學習相關 Arxiv 論文發表數量的增長態勢(藍)和摩爾定律增長率(紅)。

針對機器學習的硬件

早在 2011 年左右,谷歌研究者與系統工程師就構建了 DistBelief 這個早期分佈式系統,在 12 年早期的一些深度神經網絡中,模型使用 DistBelief 能大大提高準確性。然而,由於計算需求太大,要使用基於 CPU 的計算設備繼續提升深度模型,谷歌數據中心的計算機數量需要翻一倍。為此,這樣的現實驅動了谷歌開始考慮為神經網絡構建專用硬件,用於推理乃至訓練。

為什麼深度學習模型需要定製化硬件

深度學習模型有三大屬性,它們使得其與很多通用計算都不相同。首先,深度學習允許降低計算精度;其次,大多數模型執行的計算都是由少量運算組成,例如矩陣乘法、向量運算以及卷積運算等等;最後,過去 40 年開發了很多機制以令通用程序能在 CPU 上高效運行,但這些機制在機器學習運算中都是不必要的。

構建能進行密集低精度線性計算的硬件是非常有意義的一件事,為此,谷歌決定著手設計名為「張量處理單元」的加速器,從而加速深度學習的推理與訓練。

為了能執行推斷,Int8 的精度已經顯示出足夠強大的能力,目前研究界也在開展進一步的工作,以便採用更低的精度和更稀疏的權重來突破這一界限。TPUv1 的單個核心是由 65536 個 8 位乘加矩陣運算單元,其峰值運算量為 92 TeraOps/s。TPUv1 平均比現代 CPU 或 GPU 快 15 到 30 倍,平均能耗降低 30 到 80 倍。目前,這種芯片約佔谷歌數據中心神經網絡推理需求的 95%。

在低功率移動端設備上做推斷也非常重要,它可以極大降低網絡延遲並提升隱私保護。谷歌採用與 TPUv1 相同的設計原理,應用於更低功耗的環境而設計出 Edge TPU,它以 2W 的功耗提供了 4TOPS 的計算力。這樣的邊緣設備可以更好地在農業和日常生活中應用機器學習模型。

TPU的起源,Jeff Dean綜述後摩爾定律時代的ML硬件與算法

圖 5:谷歌張量運算單元 v2(TPUv2)的結構圖。

前面 V1 只能做推斷,但要設計定製化的機器學習訓練硬件要複雜得多。主要原因在於單塊訓練芯片無法在合理的時間內解決大部分訓練問題,因為單塊芯片無法提供足夠的算力。因此設計訓練系統實際上就是設計大規模的整體計算系統,它需要考慮各個加速芯片間的高性能互聯,從而形成及緊密耦合的超級計算機。

谷歌的第二代和第三代 TPU 旨在支持訓練和推理,TPUv2 和 TPUv3 每一個 Device 包含四塊芯片,不同的 Device 整合在一起就形成了一個整體 Pod。圖 5 展示了具有兩個核心的 TPUv2,每個核心的主要計算能力有大型的矩陣乘法單元提供,其每個時鐘週期能計算 128×128 的矩陣運算。

TPU的起源,Jeff Dean綜述後摩爾定律時代的ML硬件與算法

圖 6:谷歌 TPUv3 Pod,它包含 1024 塊 TPU 芯片。

低精度數值運算

TPUv2 和 TPUv3 使用一種名為 bfloat16 的定製化浮點運算格式,該格式與 IEEE 半精度 float 16 不同,它對機器學習更有效。自 2015 年以來,它一直是 TPU 的主要精度,英特爾 2018 年也宣佈計劃為下一代處理器增加 bfloat16 的支持。下圖展示了 IEEE fp32 單精度浮點格式、IEEE fp16 半精度浮點格式和 bfloat16 格式,它們的劃分與數值範圍都展示在內。

TPU的起源,Jeff Dean綜述後摩爾定律時代的ML硬件與算法

圖 7:三種浮點精度。

因為 bfloat16 格式的乘加器需要最少的電路,因此有可能在相同的芯片面積和功率預算中配置更多的乘法器,這能令芯片提供更強的算力與更低的能耗。與此同時,降低到 16 位精度還能減少內存的佔用與權重的體積,從而進一步提升模型的實用性。

ML 專用芯片設計的挑戰

目前機器學習研究領域正快速發展,而如今芯片設計項目通常需要 18 個月到 24 個月才能完成設計、製造與部署。為了節省成本,一般這些硬件至少需要工作三年。因此,構建 ML 硬件的計算機架構師需要預測 2 到 5 年內機器學習快速發展的領域。Jeff Dean 的經驗是,將計算機架構師、高級軟件系統構建者與機器學習研究者匯聚一堂,共同設計相關的主題,例如,「那個時期可能需要的硬件是什麼,有什麼有意思的研究趨勢,它們需要什麼樣的 ML 硬件」。

機器學習在硬件領域中的其他應用

機器學習設計芯片

機器學習另一個可以和硬件領域結合的方向是進行自動化的芯片設計,如設計定製化的 ASICs。由於芯片領域目前已有大量的人力投入,因此貿然改變整體的芯片設計是不太可能的。然而,在芯片的佈局和佈線方面機器學習是可以參與的。例如,可以將一個序列的佈局和佈線結合起來,用一個整體的評價標準——如芯片區域大小、耗時和走線長度進行評價。

如果有強化學習算法去「玩」這樣一個佈局佈線的遊戲,則可以產生很多不同的 ASIC 設計,或只有一種特定的 ASIC 設計。這隻需要設計一個獎勵函數——其中包括不同的獎勵屬性。這樣有可能能夠更快且更高效地設計出更好的佈局佈線。現在,谷歌內部已經在嘗試這些方法,而且已經有了一些初級但有效的成果。此外,針對不同的目標優化評價指標,獎勵函數中的屬性也可以進行調整,使得機器學習算法設計出滿足不同需求的芯片。

機器學習解決半導體制造問題

在半導體制造的過程中,對晶片的視覺質量檢查是可以轉移到機器學習上的,可以讓這一過程變得更自動化,或者可以提升現有方法的精確度。通過更早或更好地發現瑕疵,製造企業可以提升生產率、降低成本。

計算機系統中的學習型啟發式算法

另一個機器學習可以參與的領域則是使用學習式的啟發法,應用於計算機系統中,如編譯器、操作系統、文件系統、網絡堆棧等。計算機往往需要很多手工設計的啟發式算法用於在系統運行中尋找解決方法(如資源調用等)。

現在,機器學習也許可以替代這些算法了。基於機器學習的啟發式算法能夠考慮更多的語境信息,使得系統的運行方式更符合當前的運行模式,而不是根據一般情況設計的方法。其他機器學習可以替代的地方包括對哈希表、B 樹等傳統數據結構的替代。機器學習可以獲得這些數據的真實分佈,並使用系統進行處理,使得數據量減少、效率提升。

未來的機器學習發展

將來的機器學習模型可能和現在有一些明顯區別。將來的機器學習系統可能是在大規模機器學習加速硬件上運行的,而且單一模型可以被訓練用來完成上千甚至上百萬的任務。該模型由不同的組件和結構組成,樣本之間的數據流動可能是動態的,每種樣本都不一樣。模型可能使用類似於「稀疏權重門(sparsely-gated)」結構,混合了專家知識和學習路徑,具有很強的能力。但是對於給定的任務或樣本,模型只會激活其中一部分。

TPU的起源,Jeff Dean綜述後摩爾定律時代的ML硬件與算法

圖 8:描述這一大型、具有稀疏權重且可以完成多任務的模型。

後摩爾定律時代,簡單地壓榨硬件算力不一定能夠讓機器學習再進步了。相反,通過設計專業的機器學習硬件設備,讓機器學習解決現有的硬件設計、製造和運行方面的痛點才是出路。將來的機器學習模型可能更大,但是會朝著多任務的方向繼續發展。


分享到:


相關文章: