32分鐘訓練神經機器翻譯，速度提升45倍

專案地址：https://github.com/pytorch/fairseq

我們想讓使用者用自己喜歡的語言體驗我們的產品，同時與世界各地的人們建立聯絡。為此，我們使用神經機器學習（NMT）自動翻譯帖子和內容中的文字。我們之前關於這一課題的研究 fairseq 已經開源，這是一個序列到序列的學習庫，可供任何人訓練 NMT 模型，完成自動摘要或其他文字生成任務。

隨著 NMT 模型在從大規模單語資料（只有一種語言的資料）中的學習越來越成功，訓練速度變得越來越重要。為了適應這樣的發展，我們必須想辦法大大減少訓練時間。直到最近，NMT 模型的這種訓練在單臺機器上還需要數週，對於快速實驗和部署來說，這一速度實在太慢。

由於在精度和訓練設定上做了幾項改變，我們在短短 32 分鐘內就能訓練一個強大的 NMT 模型（參見《Scaling Neural Machine Translation》），速度是原來的 45 倍（原來需要 24 小時）。在之後的一項研究《Understanding Back-Translation at Scale》中，我們展示了這一速度大幅提升的訓練設定如何讓我們利用單語文字訓練準確得多的模型。

• 《Scaling Neural Machine Translation》：https://arxiv.org/abs/1806.00187

• 《Understanding Back-Translation at Scale》（EMNLP 2018）：https://arxiv.org/abs/1808.09381

加速訓練

我們首先在一臺配備 8 塊 Volta GPU 的英偉達 DGX-1 機器上把模型的訓練時間從近 24 小時縮減至不到 5 小時。神經網路包含數百萬個引數，這些引數在訓練過程中會被調整並通常以 32 位浮點精度儲存。然後，我們把訓練的精度從 32 位轉換為 16 位，這樣不僅能減少 GPU 記憶體佔用，還讓我們能夠使用英偉達最新 Volta GPU 架構提供的經過高度優化的 Tensor Core。降低了精度（16 位）的訓練有時會由於浮點「溢位」而導致模型質量降低。在本研究中，我們實施了一個已知的自動監控和防止溢位程式，並將訓練時間從近 24 小時減至 8.25 小時，速度提升 2.9 倍而沒有降低模型質量。

接下來，我們通過所謂的累積更新，延遲了模型更新來減少 GPU 之間的通訊。通過同步訓練模型，從而每個 GPU 可以維持一個完整、相同的模型副本，但是會處理訓練資料的不同部分。處理時需要反向傳播——在網路上被分為 forward pass 和 backward pass（前向/反向），以計算訓練模型所需的統計量。在處理每個小批量（mini-batch），GPU 會同步地相互交流結果（梯度）。這樣可能導致效率低下：首先，在 GPU 之間傳送資料會消耗時間；其次，處理較快的工作站必須等待速度較慢的工作站（即所謂散兵遊勇），然後訓練才能繼續。

後者對於文字相關模型來說是一大挑戰，不同的句子長度會使這個問題更加嚴重。我們對這兩個問題的解決方案是在工作執行緒互相通訊之前，讓每個工作執行緒處理並累積多個小批量的結果，從而延遲同步點。這將訓練時間縮減至 7.5 小時，且無損模型質量。

累積更新還能增加有效的批大小，即每一個訓練步所使用的資料量。在我們的研究設定中，批大小增大為 16 倍。這使得學習率可以翻倍，訓練時間減少到 5.2 個小時，比原始系統實現了 4.6 倍的加速。

圖：通過 (i) 減少工作站之間的通訊、(ii) 減少不同 GPU 的工作負載變化從而節約閒置時間，累加多個前向／反向步驟的梯度，實現訓練加速。

最後，我們使用 16 位浮點數來進一步提高每一個工作站的批大小，從而利用釋放出的額外 GPU 記憶體。通過將每個工作站的批大小從 3.5K 單詞提升到 5K 單詞，該研究將總訓練時間減少到 4.9 個小時，對原始系統實現了 4.9 倍的加速。

擴充套件

我們對單個機器訓練的優化還可應用於在多個機器上進行並行訓練（即分散式訓練）。我們將訓練從單個 DGX-1 機器擴充套件到 16 個機器（即 128 塊 GPU），發現可以在 37 分鐘內訓練完成同樣的模型，與原始的單機器設定相比，速度提升了 38.6 倍。

在多個機器上訓練時，另一個優化是同時發生 GPU 通訊和反向傳播。在反向傳播之後，我們將所有資訊傳達給其他工作站。典型情況下，反向傳播和通訊按順序執行，但是我們通過儘可能並行 GPU 通訊和反向傳播而節省時間。具體來說，我們在反向傳播步驟結束後立刻開始在多個工作站之間同步梯度。這樣 GPU 通訊和反向傳播發生就實現了並行，這節省了 14% 的訓練時間，且使總訓練時間從將近 24 個小時下降到 32 分鐘，速度提升了 44.7 倍。

圖：反向傳播與梯度同步的並行示意圖。

圖：每次改進後訓練總時間的變化情況。

在更多未翻譯資料上訓練

加速訓練允許我們在更大資料集上訓練模型。一般而言，訓練 NMT 模型需要原文與參考譯文對應的成對文字，例如雙語資料等。但是，雙語資料相對於網路上或其它地方常見的大量單語資料要有限地多。在第二篇論文《Understanding Back-Translation at Scale》中，我們展示瞭如何通過僅提供我們想要翻譯的語言的文字來提升模型的翻譯效能，以及利用大量的此類資料可以大幅提升模型準確度。

通過單語資料提升 NMT 模型最高效的方法之一是反向翻譯（back-translation）。如果我們的目標是訓練一個英語到德語的翻譯模型，我們可以首先訓練一個從德語翻譯到英語的模型並利用該模型翻譯所有的單語德語資料，然後在已有的和新的資料上簡單地訓練一個英語到德語的最終模型。我們的論文顯示，資料如何被翻譯是非常重要的，通過取樣故意不總是選擇最好的翻譯同樣也是非常重要的。

圖：演示反向翻譯的動畫

如果我們把 2.26 億個反向翻譯的句子新增到現有的包含 500 萬個句子的訓練資料中，我們就能大幅提升翻譯質量。下圖（左）顯示了這一系統在標準 WMT』14 英-德基準測試集上的準確率（fairseq & 單語資料）。該系統可以在 16 個 DGX-1 機器上經過 22.5 小時訓練而成。圖中還顯示了 DeepL 的準確率，這是一個依賴於高質量人工翻譯的專業翻譯服務，此前該服務在該基準測試上表現最佳。

快速翻譯

我們還提高了 fairseq 模型的翻譯速度。具體來說，我們實現了很巧妙的快取記憶體，或從計算和批處理中移除完成句子的幾個單詞而不是句子。這提高了大約 60% 的速度。下圖展示了 fairseq 和其它工具集的對比。只從 32 位浮點數變到 16 位浮點數就能提升 40% 的速度。

圖：在標準 WMT』14 英語-德語基準測試集上使用大型 Transformer 模型和 V100GPU 得到的翻譯速度。

未來計劃

擴充套件和提高自動化翻譯技術仍將是我們工作的重點。為了讓該研究專案儘快成長起來，我們共享了分散式訓練的程式碼，並將其作為我們的 fairseq 開源專案的一部分，從而使其他研究者也可以輕鬆地以更快的速度訓練 NMT 模型。進行分散式訓練僅需稍微改動訓練命令。

實現更快的 NMT 訓練是一項激動人心的發現。大幅度地降低訓練時間可以加快實驗週轉，使我們能更輕鬆地在更好的模型中做研究。把無標籤資料應用到翻譯之外的問題（如問答或自動摘要）上也是很不錯的發現。

原文連結：https://code.fb.com/ai-research/scaling-neural-machine-translation-to-bigger-data-sets-with-faster-training-and-inference/