本文最初發表於 AssemblyAI 網站,經原作者 Dylan Fox 授權,InfoQ 中文站翻譯並分享。

引   言

OpenAI 的 GPT-3 是一個令人印象深刻的深度學習模型,但是它有 1750 億個引數,相當佔用資源。儘管有不同的估計,但是這種規模的模型在一個 GPU 上的訓練需要數百年。

幸好 OpenAI 有一個由微軟提供的 NVIDIA V100 GPU 的高頻寬叢集,這讓 OpenAI 可以在幾個星期而不是幾年內就能訓練 GPT-3。這個叢集到底有多大?根據本文所述,在 1024 個 NVIDIA A100 GPU 上訓練 GPT-3 大約需要 34 天。

這個 GPU 的數量真是令人難以置信。每張 A100 GPU 的售價為 9900 美元,而我們討論的是構建這樣一個龐大的叢集需要花費將近 1000 萬美元。我們甚至還沒有考慮到電力成本,或者你實際上必須安裝 GPU 的伺服器機架,或者維護這種型別的硬體的人力成本,以及其他成本。

如今,你可以從谷歌雲這樣的公有云提供商那裡租用 A100 GPU,但按每小時 2.933908 美元計算,執行 1024 張 A100 GPU 34 天,加起來需要 2451526.58 美元。請記住,這個價格是針對單一的訓練執行的價格。

我可以繼續說下去,但問題是,訓練大型模型既昂貴又緩慢。在 AssemblyAI,我們沒有訓練 1750 億個引數範圍內的模型(謝天謝地),但是我們的語音識別模型是非常龐大的 Transformer,正在快速接近 10 億個引數。作為一家初創公司,速度和成本是我們必須不斷優化的兩件事。

這個問題的主要解決方法是在更多的 GPU 上訓練模型,但是這需要很高的成本,往往是初創公司無法承受的。近幾年來,我們學到了一些關於大型模型訓練的經驗,希望與大家分享。

模型尺寸和訓練時間

在 AssemblyAI,我們構建了大型、準確的自動語音識別(Automatic Speech Recognition,ASR)模型,並通過簡單的 語音到文字的 API 進行公開。開發人員使用我們的 API 來開發應用,來實現轉錄電話、Zoom 視訊會議、播客、視訊以及其他型別的媒體內容。

我們效能最好的自動語音識別模型是大型 Transformer,在 48 張 V100 GPU 上需要大約 3 周的時間來訓練。

32 個 NVIDIA V100s 訓練一個模型

為什麼我們這個模型的訓練需要如此長的時間和如此多的 GPU?主要原因有三個:

1. 自動語音識別模型的輸入特徵是高維、長序列

計算出每隔 10 毫秒左右的一個音訊檔案的聲譜圖,並將其作為神經網路的輸入特徵。聲譜圖的形狀 / 尺寸取決於音訊資料的取樣率,但是如果取樣率是 8000 赫茲,那麼聲譜圖中的特徵數將是 81。如果是一個 16 秒的音訊樣本,它的形狀會是 [1600, 81],這是一個相當大的特徵輸入!

下面是聲譜圖作為矩陣的一個例子:

[[[-5.7940, -5.7940, -4.1437,  ...,  0.0000,  0.0000,  0.0000],
          [-5.9598, -5.9598, -4.2630,  ...,  0.0000,  0.0000,  0.0000],
          [-5.9575, -5.9575, -4.2736,  ...,  0.0000,  0.0000,  0.0000],
          ...,
          [-4.6040, -4.6040, -3.5919,  ...,  0.0000,  0.0000,  0.0000],
          [-4.4804, -4.4804, -3.5587,  ...,  0.0000,  0.0000,  0.0000],
          [-4.4797, -4.4797, -3.6041,  ...,  0.0000,  0.0000,  0.0000]]],

        [[[-5.7940, -5.7940, -5.7940,  ...,  0.0000,  0.0000,  0.0000],
          [-5.9598, -5.9598, -5.9598,  ...,  0.0000,  0.0000,  0.0000],
          [-5.9575, -5.9575, -5.9575,  ...,  0.0000,  0.0000,  0.0000],
          ...,
          [-4.6040, -4.6040, -4.6040,  ...,  0.0000,  0.0000,  0.0000],
          [-4.4804, -4.4804, -4.4804,  ...,  0.0000,  0.0000,  0.0000],
          [-4.4797, -4.4797, -4.4797,  ...,  0.0000,  0.0000,  0.0000]]],

        [[[-5.7940, -5.7940, -5.7940,  ...,  0.0000,  0.0000,  0.0000],
          [-5.9598, -5.9598, -5.9598,  ...,  0.0000,  0.0000,  0.0000],
          [-5.9575, -5.9575, -5.9575,  ...,  0.0000,  0.0000,  0.0000],
          ...,
          [-4.6040, -4.6040, -4.6040,  ...,  0.0000,  0.0000,  0.0000],
          [-4.4804, -4.4804, -4.4804,  ...,  0.0000,  0.0000,  0.0000],
          [-4.4797, -4.4797, -4.4797,  ...,  0.0000,  0.0000,  0.0000]]]
2. 模型包含大量引數

對於基於 Transformer 的神經網路,更大的網路通常會更好。很多論文都支援這一觀點,其中 GPT-3 是最流行的例子。無論是在研究社群,還是在我們自己的內部研究中,我們都發現這種趨勢同樣適用於自動語音識別模型。

我們效能最好的模型是一個大型 Transformer,它包含近 5 億個引數。隨著引數的增加,在反向傳播過程中,梯度更新所需要的計算能力就越大。神經網路的訓練基本上可歸結為進行一堆矩陣運算。模型中的引數越多,矩陣就越大。大型矩陣需要更多的計算和 GPU 記憶體資源。

3. 對大量資料進行訓練

大型模型具有更強的建模能力,這要歸功於其引數數量的增加,為了充分利用這種建模能力,我們在近 10 萬小時的已標記的語音資料上對模型進行。舉例來說,GPT-3 是在 45TB 的文字資料上訓練的,它也可以視為 1099511626800 字左右的文字。

訓練神經網路時,需要對資料集進行多次迭代(每次迭代都被稱為“輪數”)。資料集越大,每次迭代或“輪數”的時間就越長。即使提前停止,在一個大的資料集上訓練一個大的模型,進行 20~50 次的迭代,也會花費很多時間。

如何提高迭代速度

初創公司面臨著一項艱鉅的任務:在短期內取得重大進展。被譽為“突圍型”的初創公司通常都會在最短的時間內取得最大進步。

對於一家剛起步的深度學習公司來說,這是一個艱難的挑戰。如果你的模型需要 3~4 個星期進行訓練,你是如何快速迭代的?

使用更多 GPU 訓練

減少訓練時間的最簡單方法是在更多的 GPU 上訓練模型。更多的 GPU 意味著可以使用更多的 GPU 記憶體來訓練執行。例如,假設你可以在一個 GPU 上安裝大小為 8 的 mini-batch。如果資料集中有 1000 個樣本需要迭代,這意味著需要迭代 125 個 mini-batch(每個大小為 8)。如果每次迭代需要 1 秒,那麼就需要 125 秒來迭代所有 125 個 mini-batch。

如果你有 4 個 GPU,你可以一次並行地迭代 4 個 mini-batch,而不是 1 個 mini-batch。這就是說,要完成所有 125 個 Mini-batch,只需要 32 次迭代。假定每一次迭代在 4 個 GPU 上花費 1.5 秒,這是因為 4 個 GPU 有額外的通訊開銷——然而,你仍然能夠在 48 秒內迭代完整個資料集(32*1.5)。這個速度幾乎是單個 GPU 的 3 倍。

不過,值得注意的是,更大的批量(batch)並不總是等同於更快的訓練時間。如果你的有效批量大小過大,你的模型的總體收斂性將開始受到影響。選擇適當的批量大小來訓練是你必須試驗的一項超引數,目前正針對不同的優化器(例如 LAMB 和 LARS)進行研究,這些優化器有助於緩解過大的批量大小損害收斂性的問題。

GPU 效能並非線性增長

訓練的 GPU 越多,通訊的開銷就越大。因此,在 8 個 GPU 上訓練的速度並不會比在單個 GPU 上訓練快 8 倍。在 AssemblyAI,我們使用 Horovod 來管理跨多個 GPU 上的分散式訓練執行。Horovod 是一個很棒的庫,當你在訓練叢集中增加更多的 GPU 時,它可以幫助你獲得更高的效率。

用 Horovod 的訓練時間

在測試中,我們發現 Horovod 速度大大快於 DistributedTensorFlow 和 PyTorch DistributedDataParallel。儘管如此,PyTorch 還是在積極地開發,並在快速改進。在我們的測試中,我們發現 PyTorch DistributedDataParallel 在單臺伺服器上與 Horovod 相當,但是當擴充套件訓練執行到多個伺服器時(例如,4 個伺服器,每個有 8 個 GPU),Horovod 的效能更好。

低精度訓練

大多數模型預設是使用 FP32(浮點值 32,也稱為單精度)進行訓練。使用半精度(FP16)或混合精度進行訓練,也可以加快訓練時間。

FP16 張量是 16 位,或 2 個位元組,其中每個位是 0 或 1,如 010101 10101010。FP32 張量是 32 位,或 4 位元組,如 11110000 00001111 11001100 00110011。

訓練期間更低的精度意味著更少的位元組,這意味著在訓練期間中需要的 GPU 記憶體更少,需要的頻寬也更少,而且實際硬體級操作在較新 GPU 上執行得更快,所有這些都加快了訓練速度。

使用 PyTorch,下降到 FP16 是比較容易做到的,例如 x = x.half 將一個 FP32 張量下降到 FP16。不過,要記住的是,在實踐中訓練的精確度較低,而且並不總是像在公園裡散步那麼簡單。某些操作或自定義損失函式可能不支援較低的精度,可能需要大量的超引數調整,以使你的模型在 FP16 下收斂,而且較低的精度也可能會影響模型的總體精度。

如何降低訓練成本

這很簡單:不要使用像 AWS 或谷歌雲那樣的公有云。這樣做似乎是最簡單的開始方法,但是成本會迅速增加,尤其是與下面的選擇相比。

自購硬體

如果你對管理自己的硬體感到滿意(我們不推薦這麼做),那麼購買諸如 NVIDIA TITAN X 之類的消費級 GPU 是一個比較便宜的選擇。舉例來說,每張 TITAN X 的價格大約為 3000 美元,作為消費級 GPU,其效能出乎意料的好。如果你有能力建造自己的裝置,走這條路只需支付一次硬體費用,但同時也要承擔託管和維護訓練裝置的麻煩。

一些公司如 Lambda 等,可以為你提供相對廉價的定製訓練裝置。例如,一臺配有 4 個 NVIDIA RTX A5000 和 NVLink 的機器大約需要 16500 美元。這包括記憶體、處理器、外殼等。你所要做的就是找個地方插上電源,然後支付你的電費。

專用雲服務

在 AssemblyAI,我們從 Cirrascale 租用專用伺服器。像 Cirrascale 這樣的提供商有很多,但支付專用伺服器的費用要比像 AWS 或谷歌雲這樣的大型公有云好得多。這個選擇還使你能夠自定義你所需的記憶體和處理器規格來定製你的機器,併為你選擇 GPU 提供更大的靈活性。

比如,AWS 僅提供以下 GPU:

  • NVIDIA Tesla M60 GPUs
  • NVIDIA A100
  • NVIDIA Tesla V100
  • NVIDIA K80 (these are horrible)而 Cirrascale 公司提供的 GPU 種類繁多,比如 P100s、V100s、A100s、RTX 8000s 等。

很多時候,你並不需要最昂貴的 GPU 卡(現在的 A100)來在合理的時間內訓練你的模型。而且,最新、最好的 GPU 通常不會立刻被 PyTorch 和 TensorFlow 等流行框架所支援。舉例來說,NVIDIA A100s 在得到 PyTorch 的支援前就等了一段時間。

相對於大型公有云,如 AWS 或谷歌雲,能夠根據你的訓練需求和預算定製一臺機器,對於與小型託管服務提供商合作是一個巨大的優勢。另外,由於你租用的是一臺完整的物理機器,而非 AWS/ 谷歌雲平臺那樣的虛擬化機器,因此實際的機器整體效能要好得多。

結  語

總之,訓練大型深度學習模型是許多初創公司都必需要面對的挑戰。成本可能很高,迭代時間也可能很慢,而且如果你不小心,它們會嚴重影響你的創業程式。

原文連結:

https://www.assemblyai.com/blog/how-to-train-large-deep-learning-models-as-a-startup/

來自: AI前線