百分點認知智慧實驗室出品：TPU與BERT的邂逅

百分點發表於2019-08-30

寫在前面

預訓練語言模型BERT是業界NLP領域的主流技術，百分點認知智慧實驗室在各種NLP任務中也廣泛使用BERT技術，並進行持續的優化和改進。隨著實驗室隊伍的不斷壯大，NLP任務數量的快速增加，我們在提高實驗效率並降低算力成本上開始遇到挑戰。

經過探索，我們發現使用谷歌的TPU是一種可行的解決方案。通過實驗室的研究和實踐，相比GPU，TPU能顯著提高BERT模型運算的效率，且大幅度降低所需要的成本。 

1 實驗室為什麼要採用TPU

百分點認知智慧實驗室應用預訓練語言模型的方法主要有兩種。一種方式是通過網際網路上開源的預訓練語言模型，如BERT和XLNet等，並根據上游任務建立相應的模型進行finetune。另外一種方式，在遇到對效果要求很高的NLP任務上，實驗室會根據目標語料進行預訓練語言模型的增量或重新訓練，通過新生成的預訓練模型，實驗室在文字分類、資訊抽取、問題等價性等方面均取得了重要的進展。

然而，預訓練的語言模型也帶來了一定的問題，在傳統的word embedding加LSTM或CNN為主的模型訓練當中，模型的引數很少會超過千萬級，而這樣的模型只用一塊或幾塊GPU即可在短時間內完成訓練。相比之下，base版本的BERT模型有1.5億個引數，而large版本的則超過3億，這樣的模型即使僅僅根據任務finetune也可能需要數日，而重新訓練則更是需要以月為單位計算。

計算成本開銷的增大，不但增加了應用該模型的困難，同時也為NLP演算法工程師的工作方式帶來了一定的挑戰。由於使用預訓練語言模型需要經過大量的調參，所以這使得演算法工程師的工作面臨著兩難的決策。

一方面，如果採用序列的方法進行調參，則NLP演算法工程師需要花費大量的時間進行等待，從而嚴重降低了其單位時間的產出。另一方面，如果採用並行的方式進行調參，則消耗的算力會顯著的增加。因此，在預訓練語言模型已經改變NLP領域的時代，如何探索出一種高效省錢的NLP演算法工作流程，將是NLP演算法研究員和工程師面臨的普遍問題。

值得注意的是，這個問題並非是單純的增加GPU數量就可以解決的，其原因在於在一些訓練任務中，並非增加batchsize也可以達到同樣的效果。因此在理想狀況下，NLPer需要的是在短期內可以顯著減少預訓練語言模型計算時間的高效晶片，而不僅僅是大量的簡單的堆疊現有的晶片。

為了嘗試解決這個問題，百分點認知智慧實驗室在國內首先與谷歌合作，對其提供的TPU進行比較全面的測評，並取得了積極的效果。由於TPU的中文文件和用例相對匱乏，我們首先對TPU進行簡單的介紹。

2 關於TPU的簡單介紹

TPU 即谷歌推出的 Tensor ProcessingUnit，一個只專注於神經網路計算的處理器。TPU的主要功能是矩陣運算，是一個單執行緒晶片，不需要考慮快取、分支預測、多道處理等問題的處理器，其可以在單個時鐘週期內處理數十萬次矩陣運算，而GPU在單個時鐘週期內只可以處理數百到數千的計算。

此外，GPU作為“影像處理器”，其基本的硬體架構仍然要能夠支援影像相關運算中所存在的各種需求，這意味著GPU多少受到其固有架構的桎梏，這些機制的存在使得預測GPU的處理效能和利用率變得比較困難，故在現實使用中，使GPU達到理論計算峰值是比較困難的。

而相比於GPU，TPU的實際運算表現要比GPU更加穩定，除此之外，TPU對於資料的讀取也進行了大量的優化，所以當所涉及的運算包含大量的矩陣計算時，TPU在訓練的優勢就顯現出來了。

因為BERT模型是multi-head dot-product attention，所以從數學上而言，BERT不論是前向傳播還以後向傳播計算都包含了大量矩陣運算，這使得TPU天生就非常適合計算BERT等神經網路。

當然任何理論分析都是沒有Benchmark有說服力的。那麼TPU效果究竟如何呢？下面我們來展現一下百分點認知智慧實驗室對TPU的測評結果。

3 實驗室使用TPU的結果

如上文所述，TPU在訓練BERT 模型中到底有怎樣的優勢呢？我們分別就重新預訓練語言模型和finetune語言模型的任務進行了比對。

實驗室採用了500M的文字語料，並根據主流的BERT BASE版模型的引數要求：首先將文字資料mask 十遍產生預訓練資料，然後採用了序列長度為128、12層、768維等引數進行500k步訓練。

本實驗室使用TPUV2-8進行計算並與主流GPU從運算時間和花費上進行對比，結果如下：

硬體配置

訓練步數

運算時間（天）

總花費

TPU-V2（8核心）

500k

約1.2天

$38.9

Tesla V100*8

500k

約7天

$978



實驗室分別使用了TPU和GPU對BERT模型進行500K步訓練，通過與主流的GPU Tesla V100*8 對比發現，使用GPU進行訓練花費了大約7天時間，而是用TPU進行訓練僅需要了1.2天即可完成。同時，在總費用成本上也是大量的縮減。
 
TPU在BERT預訓練模型的運算時間和總成本上完勝了當前的主流的GPU。這裡我們採用了已經訓練好的BERT模型在10萬條資料上進行十輪的finetune再次進行對比，結果如下：
 

硬體配置

運算時間（秒）

總花費

TPU-V2（8核心）

約610秒

$0.23

Tesla V100*8

約4080秒

$6.60



實驗室現在使用TPU作為加速硬體的方法對BERT模型進行十輪的fintune，如上表所示，通過和現在主流的GPUTesla V100*8進行對比發現：TPU完成fintune僅僅需要約10min的時間，而GPU完成同樣的finetune 需要超過一個小時的時間，這項技術大大提升了實驗室在NLP領域進行神經網路計算的效率，而且從總花費的角度上來看，使用TPU的成本只有GPU的3.5%左右。

因此，TPU的超高效率和低廉價格將神經網路計算變得更加“親民”了，TPU可以從根本解決了中小公司算力要求高但經費不足的顧慮，曾經那種需要幾十臺GPU幾天時間的BERT預訓練由一個TPU一天可以輕鬆解決。這讓所有的中小型企業也可以擁有之前所缺少的強大算力。
 
4 使用TPU的經驗總結

下面是實驗室對使用TPU的一個經驗總結，希望大家可以藉此排坑。

我們以租用TPU V2八核心為例，系統的說明一下建立虛擬機器例項和TPU例項的方法。進入到谷歌雲的首頁，其頁面應類似於如下：
 百分點認知智慧實驗室出品：TPU與BERT的邂逅
圖一：谷歌雲初始頁面
首先需要建立一個VM例項，在選項中進行現存、記憶體數量、系統映象等配置。

百分點認知智慧實驗室出品：TPU與BERT的邂逅

圖二：建立VM例項的介面


在該頁面的選項中，有幾個是比較重要的。

1. Machinetype。該選項決定了VM例項的執行緒數和記憶體數量的配置。一般來說，在配置系統階段，只選用最小的執行緒數和記憶體數量即可。而如果開始租用TPU，由於讀寫TensorFlowcheckpoint由CPU完成，且網路頻寬與執行緒數成正比，在TPU開始訓練後，不宜選用過小的執行緒數。

2. BootDisk。該選項指定了系統的映象。如果需要使用TPU進行計算，則要選擇支援TPU的映象。目前而言，TPU支援兩種深度學習框架，即TensorFlow和PyTorch。相比較而言，TensorFlow的支援會更為成熟，而PyTorch的支援則具有一定的實驗性。建議目前還是選用TensorFlow框架。

3. 在Identityand API access一項中，如果不存在部署的問題，建議選擇Allow full access to all Cloud APIs一項。

接下來建立TPU介面：
百分點認知智慧實驗室出品：TPU與BERT的邂逅
圖三：建立TPU介面
在建立TPU的頁面，有幾個選項值得說明：

1.       TPUtype一項中，會出現v2-8，v3-8，v3-32等選項的說明（注意不同的區域提供不同型號的TPU）。其中v2或v3為tpu的型號，-8或-32則為核心數量。最小的核心數量為8核心。在該模式下，我們可以選用搶佔式的模式。而對於大於8核心的選項，則意味著TPUpod。該模式尚不支援搶佔式（但是搶佔式正在谷歌內部進行內測）。

2.       IPaddress range一項中，如不涉及到部署，則可以填寫10.1.x.0，其中x為大於101的數字（如102，103等）。值得注意的是，如果之前已有TPU填寫了某範疇，而新建立的TPU的IP地址範疇和之也有重疊，則新建立的TPU會覆蓋掉原先的例項。

3.       Preemptibility一項為是否採用搶佔式例項的選項。如前文所述，只有V2-8和V3-8兩種型號支援建立搶佔式例項。

如果以上選項均已設定完畢，則可以點選CREATE按鈕建立TPU例項，然後就可以順利執行TPU程式了。

5 實驗室使用TPU的心得

我們在上文中介紹的VM例項和TPU例項的管理方式為眾多方法中的一種，除去以上辦法外，還可以通過命令列模型ctpuup等建立TPU例項，對此我們不再做詳細介紹。下面我們重點結合實驗室在此期間使用TPU的經驗，和大家分享心得，幫助大家排雷。

首先大家需要注意的是TPU建立完畢並開始執行，即使沒有實際的程式執行也會發生扣費。我們實驗室建議在開啟TPU之前，就先將程式碼在本地環境中調通，避免沒有必要的費用流失。從實驗室執行TPU的實戰經歷來看，我們建議使用TensorFlow的Estimator框架，因為我們只需在建立Estimator時將普通的Estimator改為TPU Estimator，即可使用TPU進行BERT神經網路預訓練。這樣大大減少了實驗室的工作量。

在進行BERT模型訓練過程中，batch_size的大小直接影響模型訓練的效能，通過對google TPU瞭解得知，每個TPU包含了8個core，建議設定batch_size大小為8的倍數。同時google 推出了TPU 的pod模式以滿足使用者對更大算力的追求，實際上就是將多個TPU進行封裝成一個整體供使用者使用，比如V3-32，就是用4個V3的TPU，因此我們建議batch_size的大小能被整體core的數目整除即可，這樣可以最大效率的利用TPU。

當然除此之外，在實戰中我們還需要處理較大的檔案在VM中的問題，因為這樣會消耗大量硬碟資源以及增加運算成本，實驗室用到了Buckets—一個價格相對親民的儲存方式來提升資源運用的效率，我們建議將較大的檔案（如BERT檔案的初始權重）儲存在bucket當中，該步驟較為簡單，故我們省略詳細介紹，有需要可以進一步閱讀https://cloud.google.com/storage/docs/creating-buckets。

6 實驗室對TPU的建議

實驗室在經過一段時間的使用後，TPU雖然在多個方面上完勝GPU，但是我們認為TPU還是有很多可改進的地方：

1.   TPU的使用門檻很高，TPU自開發以來，擁有較少的程式碼示例和文件，官方提供的例項也不夠完善，對於初學者不夠友好。尤其由於TensorFlow靜態圖的本質，這使得基於TPU的Debug比較困難；

2.   究其根本TPU是圍繞TensorFlow框架設計的硬體，實際使用過程中TPU硬體和TensorFlow版本具有較大的相互依賴性，大大減少了其可相容性，使得使用其他人工智慧框架的專案很難高效低成本地運用TPU進行運算；

3.    由於TensorFlow本身是基於靜態圖的，而TPU從本質上也只能支援靜態圖，這使得需要依賴於動態圖的應用難以在TPU上執行。這類應用包括semantic parsing和圖網絡等。

當然在TPU新技術的引進下，實驗室會迎來新的轉機和工作模式。隨著計算效率的提高，實驗室的演算法研究員可以將冗長的計算等待時間極大縮短，並且提升整體實驗室的研究效率，增強實驗專案的可控性。之前，研究員因為考慮到時間成本問題，會採用多個假設並行驗證的工作方式。因為每個實驗都有風險出現問題，並行實驗會使得研究員無法估測每個驗證的具體成功率，很有可能耗費大量算力後空手而歸。如今，研究員可以將所有假設串聯並一一快速驗證，顯著提高實驗的效率，大大降低了專案的成本風險，增加可預測性。

整體上來說，實驗室認為TPU結合BERT模型是一個不錯的開始，大大減少了我們預訓練模型的時間，顯著提升了BERT模型整體運算的效率，大幅度降低了硬體資源的算力成本。在這方面實驗室還會進一步研究和探索，讓我們期待未來更多的實踐成果。

相關文章