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

深度模型的訓練時間通常對研究者而言都是很大的挑戰，我們需要花數小時甚至數天才能知道某個小改進到底好不好。然而如果我們從一開始就考慮降低模型訓練時間，那麼很多概念都能迅速驗證。在 Facebook 開發者的試驗中，他們採用了低精度和大批量等一系列加速訓練方法，並成功地將需要 24 小時訓練的 NMT 降低到 32 分鐘。該項目相關的分散式訓練代碼已開源。

選自Facebook，作者：MICHAEL AULI、MYLE OTT、SERGEY EDUNOV，機器之心編譯。

項目地址：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 訓練是一項激動人心的發現。大幅度地降低訓練時間可以加快實驗周轉，使我們能更輕鬆地在更好的模型中做研究。把無標籤數據應用到翻譯之外的問題（如問答或自動摘要）上也是很不錯的發現。

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！