TPU和GPU的帶寬模型

前幾天，谷歌AI團(tuán)隊(duì)發(fā)布了一款新的語(yǔ)言表征模型——BERT，是來(lái)自Transformer的雙向編碼器表征。它的效果很強(qiáng)大，在11項(xiàng)NLP任務(wù)中都刷新了最佳成績(jī)。

https://arxiv.org/abs/1810.04805

在計(jì)算力方面，關(guān)于BERT選擇TPU還是GPU仍然存在爭(zhēng)議。BERT用了四天時(shí)間，在4個(gè)TPU pod上完成的（共有256個(gè)TPU芯片），這是否意味著只有谷歌才能訓(xùn)練像BERT這樣的模型呢？難道GPU已經(jīng)走到盡頭了嗎？這里需要明確兩點(diǎn)基礎(chǔ)知識(shí)：

一臺(tái)TPU是一個(gè)矩陣乘法單元，它僅可以進(jìn)行矩陣乘法和矩陣操作。在計(jì)算矩陣乘法時(shí)，它的速度很快。

進(jìn)行矩陣乘法的過(guò)程中，最慢的部分就是從主記憶體中得到元素，并將其載入處理器中。

換句話說(shuō)，矩陣相乘中，最燒錢(qián)的部分是內(nèi)存負(fù)載。對(duì)BERT來(lái)說(shuō)，矩陣相乘應(yīng)該占計(jì)算負(fù)載的90%。了解了這些背景，我們可以對(duì)這一問(wèn)題進(jìn)行小小的技術(shù)分析。

TPU和GPU的帶寬模型

TPU上的Transformer

BERT中常見(jiàn)的操作是矩陣乘法：AB=C，其中A的尺寸為256×256，B為1024×1024。TPU在對(duì)矩陣執(zhí)行相乘的過(guò)程中，會(huì)將矩陣分解成更小的128×128矩陣。這就意味著我們需要對(duì)A加載16個(gè)128×128的tile，從B中加載64個(gè)tile?？偣簿褪?664=1024個(gè)128×128的負(fù)載量。在16位的條件下，這就是32MB的數(shù)據(jù)。

現(xiàn)在我們進(jìn)一步簡(jiǎn)化它。我們假設(shè)，在進(jìn)行兩個(gè)記憶負(fù)載時(shí)沒(méi)有延遲時(shí)間，這也是有可能的，因?yàn)橥ǔＤ憧梢栽诰€程并行下隱藏內(nèi)存訪問(wèn)延遲。簡(jiǎn)單地說(shuō)，這意味著，當(dāng)我們等待一個(gè)128×128的矩陣副本完成時(shí)，已經(jīng)完成了下一個(gè)。這樣一來(lái)，我們只需要等待第一個(gè)內(nèi)存副本的完成，不用等待其他的。這就是GPU速度快的核心原因，以及為什么我們要在GPU中使用多個(gè)線程，無(wú)延遲的重疊內(nèi)存?zhèn)鬏斉c實(shí)際情況相差無(wú)幾。使用了這種簡(jiǎn)化，我們現(xiàn)在可以直接使用內(nèi)存帶寬計(jì)算為矩陣乘法加載內(nèi)存所需要的時(shí)間。如果我們查看TPU的帶寬，就會(huì)放發(fā)現(xiàn)有600GB/s，所以我們需要5.2e-05秒來(lái)傳輸32MB的數(shù)據(jù)。

GPU上的Transformer

對(duì)于GPU，過(guò)程相同，只不過(guò)使用更小的tile和更多處理器。和TPU相似，我們同時(shí)使用兩個(gè)負(fù)載來(lái)隱藏內(nèi)存延遲。對(duì)GPU來(lái)說(shuō)，16位數(shù)據(jù)的tile尺寸是96×96的。如果我們用一個(gè)V100 Tesla GPU，那么需要同時(shí)運(yùn)行160個(gè)tile，還會(huì)有稍許延遲。與TPU相比，這意味著，和兩個(gè)能處理128×128的矩陣的單元不同，GPU有160個(gè)單元（80個(gè)SM，160個(gè)線程塊，每個(gè)線程塊有兩個(gè)96×96的矩陣）。這也能保證我們可以通過(guò)并行隱藏內(nèi)存延遲。

重復(fù)以上計(jì)算過(guò)程，可以得到下面的結(jié)果：

對(duì)矩陣A，我們有33個(gè)96×96的tile；對(duì)矩陣B，我們有121個(gè)96×96的tile?？偣残枰?3*121=3993次負(fù)載，數(shù)據(jù)總量為70MB。V100每秒運(yùn)行速度為900GB，所以內(nèi)存負(fù)載可能會(huì)花7.6r-05秒。所以，我們的模型判斷，在這一場(chǎng)景下，一臺(tái)GPU比一臺(tái)TPU慢32% 。注意，對(duì)一臺(tái)RTX 2080 Ti GPU來(lái)說(shuō)，矩陣tile是一樣的，但是內(nèi)存帶寬減少到了616GB/s，說(shuō)明RTX 2080 Ti比TPU慢了54%。

注意，可用Tensor Core的TPU和GPU都能在一次運(yùn)行中分別計(jì)算矩陣乘法tile，所以，就速度來(lái)說(shuō)二者是差不多的，區(qū)別就在于內(nèi)存是如何被載入的。

在GPU上BERT的訓(xùn)練時(shí)間

利用這一數(shù)據(jù)，用V100和RTX 2080 Ti構(gòu)成的GPU群組，以及高速網(wǎng)絡(luò)和好的并行算法（例如用微軟的CNTK），我們能在64臺(tái)GPU上（相當(dāng)于四個(gè)TPU pod）、用5天多的時(shí)間或8天半的時(shí)間訓(xùn)練出BERT。在有八臺(tái)GPU的設(shè)備上，使用任意軟件和并行算法（PyTorch或者TensorFlow），我們訓(xùn)練BERT需要42天或者68天。對(duì)于標(biāo)準(zhǔn)的四個(gè)GPU的臺(tái)式機(jī)，我們需要99天。

帶寬模型的限制

帶寬模型最大的限制就是這些計(jì)算是針對(duì)特定矩陣大小的，計(jì)算的難度在各種尺寸之間都不同。例如，如果你的batch size是128，那么GPU的速度會(huì)比TPU稍快一點(diǎn)。如果batch size小于128，GPU的速度就會(huì)快很多。增加矩陣B的尺寸會(huì)讓TPU比GPU快得多。在BERT的原始論文中，研究人員是針對(duì)TPU進(jìn)行矩陣尺寸的調(diào)整，如果你用GPU訓(xùn)練的話就不要參考了。

未來(lái)可能遇到的限制包括融合運(yùn)算，TPU可以計(jì)算額外的操作，例如非線性激活函數(shù)或矩陣乘法中的偏差。這意味著，TPU不需要從較慢的全局內(nèi)存中進(jìn)行加載。另外，GPU也支持這種操作，但是英偉達(dá)還未將它們實(shí)現(xiàn)，所以GPU用戶可能無(wú)法使用這一方法。所以，用戶可能會(huì)遇到1.6%的速度變慢，例如，如果你應(yīng)用一個(gè)非線性函數(shù)和一個(gè)偏見(jiàn)，那么TPU可能會(huì)比GPU快3.2%。

32位、16位和8位的區(qū)別

如果在32位情況下重復(fù)上面的計(jì)算過(guò)程，那么TPU會(huì)比之前快5.3倍。所以數(shù)據(jù)類型的尺寸在TPU和GPU的選擇上還是很重要的。

TPU不支持8位訓(xùn)練，但是圖靈GPU可以。所以，我們也可以看看8位矩陣乘法會(huì)有怎樣的表現(xiàn)。我此前總結(jié)了8位模型的表現(xiàn)，發(fā)現(xiàn)并不難訓(xùn)練。如果我們重復(fù)上面的訓(xùn)練，用8位GPU訓(xùn)練，就會(huì)發(fā)現(xiàn)GPU比TPU快了3倍。在四臺(tái)RTX 2080 Ti上的8位訓(xùn)練需要21天。

結(jié)語(yǔ)

TPU在訓(xùn)練類似BERT的模型時(shí)，比GPU快32%到54%。你也可以用40到70天的時(shí)間，在八臺(tái)GPU上復(fù)現(xiàn)BERT。而用四臺(tái)普通GPU，在16位的情況下訓(xùn)練出BERT需要99天，在8位情況下則需要21天。