機(jī)器學(xué)習(xí)如何影響計(jì)算機(jī)硬件設(shè)計(jì)1

前 言

為什么芯片設(shè)計(jì)需要很長(zhǎng)時(shí)間？能不能加速芯片設(shè)計(jì)周期？能否在幾天或幾周之內(nèi)完成芯片的設(shè)計(jì)？這是一個(gè)非常有野心的目標(biāo)。過(guò)去十年，機(jī)器學(xué)習(xí)的發(fā)展離不開(kāi)系統(tǒng)和硬件的進(jìn)步，現(xiàn)在機(jī)器學(xué)習(xí)正在促使系統(tǒng)和硬件發(fā)生變革。

Google在這個(gè)領(lǐng)域已率先出發(fā)。在第58屆DAC大會(huì)上，Google AI負(fù)責(zé)人Jeff Dean分享了《機(jī)器學(xué)習(xí)在硬件設(shè)計(jì)中的潛力》，他介紹了神經(jīng)網(wǎng)絡(luò)發(fā)展的黃金十年，機(jī)器學(xué)習(xí)如何影響計(jì)算機(jī)硬件設(shè)計(jì)以及如何通過(guò)機(jī)器學(xué)習(xí)解決硬件設(shè)計(jì)中的難題，并展望了硬件設(shè)計(jì)的發(fā)展方向。

他的演講重點(diǎn)在于Google如何使用機(jī)器學(xué)習(xí)優(yōu)化芯片設(shè)計(jì)流程，這主要包括架構(gòu)搜索和RTL綜合、驗(yàn)證、布局與布線（Placement and routing）三大階段。在架構(gòu)搜索階段，Google提出了FAST架構(gòu)自動(dòng)優(yōu)化硬件加速器的設(shè)計(jì)，而在驗(yàn)證階段，他們認(rèn)為使用深度表示學(xué)習(xí)可提升驗(yàn)證效率，在布局與布線階段，則主要采用了強(qiáng)化學(xué)習(xí)技術(shù)進(jìn)行優(yōu)化。

以下是他的演講內(nèi)容：

1

神經(jīng)網(wǎng)絡(luò)的黃金十年

制造出像人一樣智能的計(jì)算機(jī)一直是人工智能研究人員的夢(mèng)想。而機(jī)器學(xué)習(xí)是人工智能研究的一個(gè)子集，它正在取得很多進(jìn)步?，F(xiàn)在大家普遍認(rèn)為，通過(guò)編程讓計(jì)算機(jī)變得“聰明”到能觀察世界并理解其含義，比直接將大量知識(shí)手動(dòng)編碼到人工智能系統(tǒng)中更容易。

神經(jīng)網(wǎng)絡(luò)技術(shù)是一種非常重要的機(jī)器學(xué)習(xí)技術(shù)。神經(jīng)網(wǎng)絡(luò)一詞出現(xiàn)于1980年代左右，是計(jì)算機(jī)科學(xué)術(shù)語(yǔ)中一個(gè)相當(dāng)古老的概念。雖然它當(dāng)時(shí)并沒(méi)有真正產(chǎn)生巨大的影響，但有些人堅(jiān)信這是正確的抽象。

本科時(shí)，我寫(xiě)了一篇關(guān)于神經(jīng)網(wǎng)絡(luò)并行訓(xùn)練的論文，我認(rèn)為如果可以使用64個(gè)處理器而不是一個(gè)處理器來(lái)訓(xùn)練神經(jīng)網(wǎng)絡(luò)，那就太棒了。然而事實(shí)證明，我們需要大約100萬(wàn)倍的算力才能讓它真正做好工作。

2009年前后，神經(jīng)網(wǎng)絡(luò)技術(shù)逐漸火熱起來(lái)，因?yàn)槲覀冮_(kāi)始有了足夠的算力讓它變得有效，以解決現(xiàn)實(shí)世界的問(wèn)題以及我們不知道如何解決的其他問(wèn)題。2010年代至今是機(jī)器學(xué)習(xí)取得顯著進(jìn)步的十年。

是什么導(dǎo)致了神經(jīng)網(wǎng)絡(luò)技術(shù)的變革？我們現(xiàn)在正在做的很多工作與1980年代的通用算法差不多，但我們擁有越來(lái)越多的新模型、新優(yōu)化方法等，因此可以更好地工作，并且我們有更多的算力，可以在更多數(shù)據(jù)上訓(xùn)練這些模型，支撐我們使用更大型的模型來(lái)更好地解決問(wèn)題。

在探討設(shè)計(jì)自動(dòng)化方面之前，我們先來(lái)看看一些真實(shí)世界的例子。首先是語(yǔ)音識(shí)別。在使用深度學(xué)習(xí)方法之前，語(yǔ)音識(shí)別很難得到實(shí)際應(yīng)用。但隨后，使用機(jī)器學(xué)習(xí)和神經(jīng)網(wǎng)絡(luò)技術(shù)，大幅降低了詞語(yǔ)的識(shí)別錯(cuò)誤率。

幾年后，我們將錯(cuò)誤率降低到5%左右，讓語(yǔ)音識(shí)別更加實(shí)用，而現(xiàn)在，在不聯(lián)網(wǎng)的設(shè)備里，我們都可以做到僅僅4%左右的錯(cuò)誤率。這樣的模型被部署在人們的手機(jī)里面，隨時(shí)隨地幫助人們識(shí)別自己的語(yǔ)音。

計(jì)算機(jī)視覺(jué)方面也取得了巨大的進(jìn)步。2012年左右，Alex Krizhevsky、Ilya Sutskever和Geoffrey Hinton在ImageNet比賽中首次使用了AlexNet，錯(cuò)誤率得到顯著降低，并在當(dāng)年奪得桂冠。

后一年的ImageNet比賽中，幾乎所有參賽者都使用深度學(xué)習(xí)方法，研究人員則進(jìn)一步放棄了傳統(tǒng)的方法。其中，2015年，由何愷明等微軟研究人員提出ResNet更進(jìn)一步降低了錯(cuò)誤率。

當(dāng)時(shí)的斯坦福大學(xué)研究生Andrej Karpathy正在幫助運(yùn)營(yíng)ImageNet比賽，他想知道如果人工識(shí)別這項(xiàng)艱難的任務(wù)，錯(cuò)誤率會(huì)是多少。在上千個(gè)類別中有40種狗，你必須能夠看著一張照片說(shuō)：“哦，那是一只羅威納犬，不是一只大力金剛?cè)?，或者其他品種的狗。” 經(jīng)過(guò)一百個(gè)小時(shí)的訓(xùn)練，他將錯(cuò)誤率降到了5%。

這是一項(xiàng)非常艱難的任務(wù)，將計(jì)算機(jī)識(shí)別錯(cuò)誤率從2011年的26%降低到2017年的2%是一件很了不起的事，過(guò)去計(jì)算機(jī)無(wú)法識(shí)別的東西，現(xiàn)在已經(jīng)可以識(shí)別。自然語(yǔ)言處理、機(jī)器翻譯和語(yǔ)言理解中也經(jīng)歷了類似的故事。

此外，開(kāi)源框架確實(shí)使世界各地的許多人能夠應(yīng)用機(jī)器學(xué)習(xí)技術(shù)，TensorFlow就是其中之一。

大約在2015年11月，我們開(kāi)源了TensorFlow以及供Google內(nèi)部使用的工具。TensorFlow對(duì)世界產(chǎn)生了相當(dāng)大的影響，它已經(jīng)被下載了大約5000萬(wàn)次，當(dāng)然也出現(xiàn)了很多其他框架，比如JAX、PyTorch等等。

世界各地的人們能夠?qū)C(jī)器學(xué)習(xí)用于各種了不起的用途，例如醫(yī)療保健、機(jī)器人技術(shù)、自動(dòng)駕駛等等，這些領(lǐng)域都是通過(guò)機(jī)器學(xué)習(xí)方法來(lái)理解周圍的世界，進(jìn)而推動(dòng)領(lǐng)域的發(fā)展。

2

機(jī)器學(xué)習(xí)改變計(jì)算機(jī)設(shè)計(jì)方式

ML研究社區(qū)中的許多成功源自使用更多算力和更大的模型，更多的算力促進(jìn)了機(jī)器學(xué)習(xí)研究領(lǐng)域中重要成果的產(chǎn)生。深度學(xué)習(xí)的發(fā)展正在深刻改變計(jì)算機(jī)的結(jié)構(gòu)?，F(xiàn)在，我們想圍繞機(jī)器學(xué)習(xí)計(jì)算類型構(gòu)建專門(mén)的計(jì)算機(jī)。

近年來(lái)，我們已經(jīng)在Google做了很多類似的工作，其中TPU（張量處理單元）是我們構(gòu)建定制處理器的一種方法，這些處理器專為神經(jīng)網(wǎng)絡(luò)和機(jī)器學(xué)習(xí)模型而設(shè)計(jì)。

TPU v1是我們第一個(gè)針對(duì)推理的產(chǎn)品，當(dāng)你擁有經(jīng)過(guò)訓(xùn)練的模型，并且只想獲得已投入生產(chǎn)使用的模型的預(yù)測(cè)結(jié)果，那它就很適合，它已經(jīng)被用于神經(jīng)機(jī)器翻譯的搜索查詢、AlphaGo比賽等應(yīng)用中。

后來(lái)我們還構(gòu)建了一系列處理器。TPU v2旨在連接在一起形成稱為Pod的強(qiáng)大配置，因此其中的256個(gè)加速器芯片通過(guò)高速互聯(lián)緊緊連接在一起。TPU v3則增加了水冷裝置。

TPU v4 Pod不僅可以達(dá)到ExaFLOP級(jí)的算力，它還讓我們能夠在更大的模型訓(xùn)練中達(dá)到SOTA效果，并嘗試做更多的事情。

以ResNet-50模型為例，在8塊P100 GPU上訓(xùn)練完ResNet-50需要29小時(shí)，而在2021年6月的MLPerf競(jìng)賽中，TPU v4 pod僅耗時(shí)14秒就完成了訓(xùn)練。但我們的目的不僅僅是在14秒內(nèi)訓(xùn)練完ResNet，而是想把這種強(qiáng)大的算力用于訓(xùn)練其他更先進(jìn)的模型。

可以看到，從一開(kāi)始的29小時(shí)到后來(lái)的14秒，模型的訓(xùn)練速度提高了7500倍。我認(rèn)為實(shí)現(xiàn)快速迭代對(duì)于機(jī)器學(xué)習(xí)非常重要，這樣才能方便研究者試驗(yàn)不同想法。

基于機(jī)器學(xué)習(xí)的計(jì)算方式越來(lái)越重要，計(jì)算機(jī)也正在往更適應(yīng)機(jī)器學(xué)習(xí)計(jì)算方式的方向上演進(jìn)。但深度學(xué)習(xí)有可能影響計(jì)算機(jī)的設(shè)計(jì)方式嗎？我認(rèn)為，答案是肯定的。

3

機(jī)器學(xué)習(xí)縮短芯片設(shè)計(jì)周期

目前，芯片的設(shè)計(jì)周期非常長(zhǎng)，需要幾十甚至幾百人的專業(yè)團(tuán)隊(duì)花費(fèi)數(shù)年的努力。從構(gòu)思到完成設(shè)計(jì)，再到成功生產(chǎn)，中間的時(shí)間間隔十分漫長(zhǎng)。但如果將來(lái)設(shè)計(jì)芯片只需要幾個(gè)人花費(fèi)幾周時(shí)間呢？這是一個(gè)非常理想的愿景，也是研發(fā)人員當(dāng)前的目標(biāo)。

如上圖所示，芯片設(shè)計(jì)包含四個(gè)階段：架構(gòu)探索→RTL綜合→驗(yàn)證→布局和布線。完成設(shè)計(jì)之后，在制作生產(chǎn)環(huán)節(jié)需要進(jìn)行布局和布線（Placement & Routing），有沒(méi)有更快、更高質(zhì)量的布局和布線方法？驗(yàn)證是非常耗時(shí)的一步，能不能用更少的測(cè)試次數(shù)涵蓋更多的測(cè)試項(xiàng)目？有沒(méi)有自動(dòng)進(jìn)行架構(gòu)探索和RTL綜合的方法？目前，我們的芯片架構(gòu)探索只針對(duì)幾種重要的應(yīng)用，但我們終將要把目光擴(kuò)大。

布局與布線

首先，關(guān)于布局和布線，Google在2020年4月發(fā)表過(guò)一篇論文Chip Placement with Deep Reinforcement Learning，2021年6月又在Nature上發(fā)表了A graph placement methodology for fast chip design。

我們知道強(qiáng)化學(xué)習(xí)的大致原理：機(jī)器執(zhí)行某些決定，然后接收獎(jiǎng)勵(lì)（reward）信號(hào)，了解這些決定帶來(lái)什么結(jié)果，再據(jù)此調(diào)整下一步?jīng)Q定。

因此，強(qiáng)化學(xué)習(xí)非常適合棋類游戲，比如國(guó)際象棋和圍棋。棋類游戲有明確的輸贏結(jié)果，機(jī)器下一盤(pán)棋，總共有50到100次走棋，機(jī)器可以根據(jù)最終的輸贏結(jié)果評(píng)定自己和對(duì)手的整套走棋方法的有效性，從而不斷調(diào)整自己的走棋，提高下棋水平。

那么ASIC芯片布局這項(xiàng)任務(wù)能不能也由強(qiáng)化學(xué)習(xí)智能體來(lái)完成呢？

這個(gè)問(wèn)題有三個(gè)難點(diǎn)。第一，芯片布局比圍棋復(fù)雜得多，圍棋有10^{360}種可能情況，芯片布局卻有10^{9000}種。

第二，圍棋只有“贏”這一個(gè)目標(biāo)，但芯片布局有多個(gè)目標(biāo)，需要權(quán)衡芯片面積、時(shí)序、擁塞、設(shè)計(jì)規(guī)則等問(wèn)題，以找到最佳方案。

第三，使用真實(shí)獎(jiǎng)勵(lì)函數(shù)（true reward function）來(lái)評(píng)估效果的成本非常高。當(dāng)智能體執(zhí)行了某種芯片布局方案后，就需要判斷這個(gè)方案好不好。如果使用EDA工具，每次迭代都要花上很多個(gè)小時(shí)，我們希望將每次迭代所需時(shí)間縮減為50微秒或50毫秒。

利用強(qiáng)化學(xué)習(xí)進(jìn)行芯片布局的步驟如下：首先從空白底座開(kāi)始，運(yùn)用分布式PPO算法（強(qiáng)化學(xué)習(xí)的常用算法）進(jìn)行設(shè)計(jì)，然后完成每個(gè)節(jié)點(diǎn)的布局放置，最后進(jìn)行評(píng)估。

評(píng)估步驟使用的是代理獎(jiǎng)勵(lì)函數(shù)（proxy reward function），效果和真實(shí)獎(jiǎng)勵(lì)函數(shù)相近，但成本低得多。在一秒或半秒內(nèi)就可以完成對(duì)本次布局方案的評(píng)估，然后指出可優(yōu)化之處。

構(gòu)建獎(jiǎng)勵(lì)函數(shù)需要結(jié)合多個(gè)不同的目標(biāo)函數(shù)，例如線長(zhǎng)、擁塞和密度，并分別為這些目標(biāo)函數(shù)設(shè)定權(quán)重。

如上圖所示，布局優(yōu)化采取的是混合方式。強(qiáng)化學(xué)習(xí)智能體每次放置宏（macro），然后通過(guò)力導(dǎo)向方法（force-directed method）放置標(biāo)準(zhǔn)單元。

上圖來(lái)自前面提到的Nature論文，展示了更多芯片架構(gòu)的細(xì)節(jié)。

上圖展示了一個(gè)TPU設(shè)計(jì)塊的布局與布線結(jié)果。白色區(qū)域是宏，綠色區(qū)域是標(biāo)準(zhǔn)單元群（standard cell clusters）。

圖中左邊是人類專家完成的設(shè)計(jì)，從中可以看出一些規(guī)律。人類專家傾向于把宏沿邊緣放置，把標(biāo)準(zhǔn)單元放在中間。一名人類專家需要6~8周完成這個(gè)布局，線長(zhǎng)為57.07米。圖中右邊是由智能體（ML placer）完成的布局，耗時(shí)24小時(shí)，線長(zhǎng)55.42米，違反設(shè)計(jì)規(guī)則的地方比人類專家略多，但幾乎可以忽略。