谷歌利用 AI 設計晶元

就好像嫌摩爾定律在晶體管密度和成本方面改進的步伐在放緩還不夠糟糕，設計晶元的成本以及蝕刻晶元的工廠的成本也在上升。為了保持IT創新闊步前行，任何這些成本得到節省都將大受歡迎。

目前，晶元設計界大有希望的研究領域之一是，使用機器學習技術來切實幫助處理設計過程中的一些任務。晶元設計中使用機器學習也是谷歌研究部門的高級研究員Jeff Dean在本周於舊金山舉行的2020年國際固態電路大會上發表主題演講時暢談的話題之一，他幫助谷歌這家超大規模提供商發明了許多關鍵技術。

事實證明，谷歌對計算引擎絕非一時心血來潮。谷歌是全球最大的CPU和GPU消費者之一，也是TPU的設計者，TPU的應用範圍從邊緣到數據中心，不一而足，既可用於機器學習推理又可以用於機器學習訓練。因此，這不僅僅是這個搜索引擎巨頭和公共雲提供商的科研活動——如果它打算繼續推進TPU路線圖，如果它像競爭對手AWS那樣決定開始設計自己的定製Arm伺服器晶元，或決定為其手機及其他消費級設備設計定製Arm晶元，更是具有現實意義。

出於一定的偶然性，谷歌為了跨大量不同類型的計算引擎一直運行機器學習模型所做的一些工作回過頭來進入到它為了使ASIC上IP功能塊（IP block）的一些放置和布線實現自動化而在做的工作。

雖然谷歌早在2018年5月展示的TPUv3系統的pod可以將1024個張量處理器（核心數量增加一倍，時鐘頻率提高15%左右）組合起來，以提供每秒106千萬億次浮點運算的16位半精度乘法總性能（帶32位累加），使用谷歌自家的非常巧妙的bfloat16數據格式。這些TPUv3晶元都使用32×32環形網格實現交叉耦合，那樣它們就能共享數據，而每個TPUv3核心都有自己的HBM2存儲器模組。這個TPUv3 pod是巨大的計算聚合體，能進行機器學習訓練或機器學習推理，但是未必如谷歌需要製造的那麼大。

就一句話，為了規避風險，谷歌採用結合CPU和GPU的混合架構（也許有一天會結合其他加速器）用於處理強化學習工作負載，因此Dean及其在谷歌的同事從事的研究成果也應用於ASIC設計上。

Dean解釋道：「趨勢之一就是模型越來越大。因此，整個模型在單單一塊晶元上不一定裝得下。如果你有龐大模型，那麼模型並行處理（將模型分散於多塊晶元之間）很重要，通過為它提供一堆計算設備來獲得良好性能並非易事，如何有效地處理這件事也並不明顯。」

由於AI系統具有異構性，這不像拿來用於在大規模並行超級計算機上分派工作的消息傳遞介面（MPI），稍加改動後部署到TensorFlow之類的機器學習框架上那麼簡單。但這可能是將機器學習訓練工作負載分布在眾多計算單元上的一條思路，已有人這麼做了。與其他超大規模提供商一樣，谷歌往往構建自己的框架、協議和數據存儲系統，當然依託其他技術。

由於神經網路模型變得比存儲器空間還要龐大，超過了單單一個CPU、GPU或TPU的計算能力，設備放置顯得尤為重要；設備放置是指，在合適的時間將合適的神經網路（或體現神經網路的代碼的一部分）放在合適的設備上，以便在整個應用中獲得最大吞吐量。而這個問題迅速變得嚴重起來，框架和硬體跟不上步伐。不妨看一看：

參數的數量就在不斷增加，並行使用的設備數量也在不斷增加。實際上，讓128個GPU或128個TPUv3處理器（這就是為什麼上圖中有512個核心）協同工作是相當了不起的成就，可與早期超級計算機最強大的處理能力相提並論。20多年前，使用MPI的鬆散耦合的大規模並行超級計算機唱主角，搭載物理共享存儲器的聯合NUMA伺服器成為了HPC的標準。隨著越來越多的設備以某種方式組合起來以處理這些模型，谷歌一直在嘗試使用強化學習（機器學習的一個特殊分支），搞清楚全體模型在一組CPU和GPU上運行時，在任何特定的時間哪裡運行神經網路模型效果最好。在這種情況下，制定初始策略以調派處理任務的神經網路模型，然後結果反饋回到模型中作進一步調整，以便那些模型運行起來越來越高效。

2017年，谷歌訓練強化學習模型以完成這項工作（論文詳見https://arxiv.org/abs/1706.04972），這是編碼器和解碼器最終放置的樣子；而與專家手動放置的神經網路相比，機器學習模型將工作負載放置在接受測試的系統中的兩個CPU和四個GPU上，其最終運行時間縮短19.3%。Dean補充道，利用強化學習將神經網路工作負載放置在計算引擎上「做了有點不合常理的事」才獲得這種結果，許多機器學習應用似乎是這種情況。問題在於，它無法獲得大量的強化學習計算能力將工作負載放置在設備上，以運行自行訓練的神經網路。2018年，谷歌開展了研究以表明如何將計算圖擴展到80000多個操作（節點），去年穀歌為有50000多個操作（節點）的數據流圖創建了所謂的廣義設備放置方案。

Dean說：「隨後我們開始考慮使用該方法，而不是用它將軟體計算放置在不同的計算設備上，我們開始考慮將它用於ASIC晶元設計中的放置和布線，因為問題看起來有點相似。強化學習確實適用於處理規則清楚的棘手問題（比如國際象棋或圍棋）；實際上我們開始問自己：我們能否獲得強化學習模型成功地做好ASIC晶元布局這項工作？」

據Dean聲稱，這麼做有幾個挑戰。舉例說，國際象棋和圍棋都有一個目標，那就是贏得比賽而不是輸掉比賽。（它們是同一枚硬幣的兩面。）就ASIC上放置IP功能塊和IP功能塊之間的布線而言，沒有簡單的贏或輸，你要關注很多目標，比如面積、時間、擁塞和設計規則等。更困難的是，神經網路模型為IP功能塊放置所要管理的潛在狀態數量眾多，如下圖所示：

最後，決定IP功能塊放置的實際獎勵函數（在EDA工具中運行）需要花大量時間來運行。

Dean解釋：「於是我們有了一種架構（我不會透露很多細節），但實際上它設法把構成晶元設計的一堆東西放到晶圓上。」他展示了將IP功能塊放置在低功率機器學習加速器晶元上的一些結果（我們假設該晶元是谷歌為其智能手機製造的邊緣TPU），一些地方有意模糊處理，那樣外人無法了解晶元的細節。「我們有一隊專家來放置該IP功能塊，他們有幾個代理獎勵函數，我們評估這些函數成本很低；我們在兩秒鐘而不是幾小時內對它們進行了評估，這點非常重要，因為強化學習中要迭代多次。因此，我們有一個基於機器學習的放置系統；你可以看到，它稍微擴展了邏輯，而不是讓邏輯分布在一個呈矩形的地方，這使它能夠在擁塞和導線長度方面都有所改進。我們在迄今為止試用的所有不同IP功能塊上取得了與人相當或更勝一籌的結果。」

不管怎樣，這是針對強化學習網路的低功耗機器學習加速器與放置IP功能塊的人員的比較：

而這張表格顯示了手動放置和布線與藉助機器學習自動放置和布線之間的區別：

最後，這表明了與人相比，強化學習網路如何處理TPU晶元上的IP功能塊：

看看這些AI設計的IP功能塊與人類設計的笛卡爾功能塊相比是多麼地合乎自然。太棒了。

谷歌完成這項工作後問了這個問題：我們可以訓練快速有效地放置以前從未見過的新設計的普通代理嗎？這正是你在設計新晶元時關注的重點。於是，谷歌針對來自TPU架構的四個不同IP功能塊測試了這個廣義模型，然後還在「Ariane」RISC-V處理器架構上進行了測試。下比是數據對比：

以下是關於在Ariane RISC-V晶元上進行放置和布線的更多數據：

Dean說：「你可以看到其他設計方面的情況實際上顯著改善了結果，因此從本質上講，你可以在12個小時內得到深藍色條，」他指上面第一張圖，然後以上面第二張圖繼續說事。「這張圖顯示了導線成本，我們可以看到如果你從頭開始訓練，實際上它花一點時間才有某種突破性的真知灼見，才能顯著降低導線成本，而事先訓練的策略只要看看其他設計，很快就能對晶元設計有一番大致的了解。」

就像我們進行大量模擬以便更準確預報天氣一樣，Dean說晶元設計中這種基於AI的IP功能塊放置和布線可用於快速生成許多不同的布局，自有不同的美中不足。而萬一需要添加某個功能，基於AI的晶元設計可以快速重新進行布局，不用花好幾個月來做。

而最重要的是，這種自動化設計輔助有望從根本上降低設計新晶元的成本。這類成本在急劇增長，而我們看到的數據（感謝IT行業知名人士、Arista Networks主席兼首席技術官Andy Bechtolsheim），採用16納米工藝的先進晶元設計平均成本高達1.063億美元，改而10納米會使成本增加至1.744億美元，改而7納米需要2.978億美元的成本，預計5納米晶元的成本在5.422億美元左右。其中近一半的成本用於軟體，現在依然如此。因此，我們知道竭力消除哪方面的部分成本，而機器學習可助一臂之力。

問題是，晶元設計軟體開發商會不會融入AI，並促使晶元設計領域迎來可以真正稱為寒武紀的爆炸式增長，然後實現批量生產？這方面會發生什麼、谷歌開展的此類研究將如何提供幫助值得拭目以待。

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！