前幾天，谷歌AI團隊發(fā)布了一款新的語言表征模型——BERT，是來自Transformer的雙向編碼器表征。它的效果很強大，在11項NLP任務(wù)中都刷新了最佳成績。

https://arxiv.org/abs/1810.04805

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

一臺TPU是一個矩陣乘法單元，它僅可以進行矩陣乘法和矩陣操作。在計算矩陣乘法時，它的速度很快。

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

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

TPU和GPU的帶寬模型

TPU上的Transformer

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

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

GPU上的Transformer

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

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

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

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

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

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

帶寬模型的限制

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

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

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

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

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

結(jié)語

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