IBM全新機器學習庫SnapML：訓練提速46倍，40億樣本僅需91.5秒

策劃編輯 | Natalie

譯者 | 無明

編輯 | Emily

AI 前線導讀：在拉斯維加斯舉行的 IBM THINK 2018 大會上，IBM 宣布了一項與機器學習性能有關的重大突破。他們使用一組由 Criteo 實驗室發布的廣告數據集（包含了 40 億個訓練樣本）來訓練邏輯回歸分類器，結果只花了 91.5 秒就完成了整個訓練過程，這一成績比之前報告的最好成績快了 46 倍。2017 年 2 月，谷歌雲官方博客披露了一組實驗數據，他們在谷歌雲平台上使用 TensorFlow 訓練同樣的模型，結果花了 70 分鐘。

Snap ML刷新了已有的基準性能數據。

更多乾貨內容請關注微信公眾號「AI 前線」，（ID：ai-front）

位於蘇黎世 IBM 研究室的一個團隊在過去兩年開發了一個新的機器學習庫，叫作 Snap ML（IBM Snap Machine Learning）。之所以把它叫作 Snap，是因為他們相信用它來訓練模型也就是彈指一揮間的事。這個機器學習庫提供了非常快的訓練速度，可以在現代 CPU 或 GPU 系統上訓練主流的機器學習模型，或者在有新數據可用時保證已有模型的速度保持在線速水平（也就是網路能夠支持的最快速度），這意味著更低的計算成本、更少的能源消耗和更快的開發速度。

與時間賽跑

數據的增長為機器學習和人工智慧的大規模應用奠定了基礎。大數據集為訓練出豐富的模型提供了更多可能性，但當訓練樣本或特徵達到數十億的規模時，即使是訓練最簡單的模型也需要花費很長時間。在為關鍵應用（如天氣預報和金融欺詐檢測）研究、開發和部署大規模的機器學習模型時，過長的時間周期將成為嚴重障礙。

Snap ML 適用於多種場景。首先是擁有大量數據的應用程序，在這種場景下，需要解決訓練時間瓶頸問題。其次是實時或近實時的應用程序，在這種場景下，模型會快速發生變更，同樣對訓練速度提出了很高的要求。第三個是集成學習（Ensemble Learning）領域。眾所周知，現如今要想在數據科學領域的競爭中勝出，很大程度上要依賴大型的模型群組。為了設計出有競爭力的模型群組，數據科學家通常需要花費大量的時間去嘗試各種模型組合和大量的超參數。在這種情況下，只有把訓練速度提升一個數量級，才有可能加快研發速度，以便在競爭中勝出。

機器學習的效率、結果和洞見對於大大小小的企業來說都至關重要，而機器學習需要消耗大量的計算資源。計算資源的成本在持續攀升，所以對於企業來說，時間就是金錢。

Snap ML 採用最先進的演算法和最優的系統設計，充分利用硬體資源，從而提供了更快的訓練速度。Snap ML 的三個主要特性如下：

分散式訓練：Snap ML 是一個可以自由伸縮的數據並行框架，訓練數據集可以超出單台機器的內存容量，這對於大型應用程序來說至關重要。

GPU 加速：Snap ML 實現了一些特殊的程序，可以充分利用 GPU 的並行特性，同時在 GPU 內存中保留數據，減少數據傳輸開銷。

稀疏數據結構：大部分機器學習數據集都是稀疏的，為此，他們對稀疏數據結構的相關演算法進行了優化。

萬億級別的基準測試

由 Criteo 實驗室發布的廣告點擊日誌訓練數據集包含了 40 億個訓練樣本，每個樣本都有一個「標籤」和一組相關的匿名特徵，其中標籤用於標識樣本是否由用戶點擊在線廣告所生成。基於這組數據集進行機器學習，可以預測出新用戶是否有可能點擊某個廣告。這個數據集是目前公開發布的最大的可用數據集之一，帶有 24 天廣告數據，平均每天包含 1 億 6 千萬個訓練樣本。

Snap ML 團隊將 Snap ML 部署在 IBM Power System AC922 伺服器上，每台伺服器配備了 4 個 NVIDIA Tesla V100 GPU 和兩個 Power9 CPU，它們通過 NVLINK 2.0 介面與主機連接，而伺服器之間則通過無限帶寬網路進行通信。他們在這些伺服器上訓練了一個邏輯回歸分類器，時間為 91.5 秒，損失為 0.1292。

2017 年 2 月，谷歌花了 70 分鐘在他們的雲平台上使用 TensorFlow 訓練同樣的數據集。而且，他們使用了 60 台工作機和 29 台參數機。相比 TensorFlow，Snap ML 獲得相同的損失，但速度卻快了 46 倍。

在訓練過程中需要解決的一個主要技術問題是，如何將那麼大的數據集塞進 GPU 內存中。因為數據集太大，GPU 內存又很小，在訓練過程中就需要不斷地將數據移進移出 GPU 內存。為了分析 GPU 內核和數據拷貝分別花了多少時間，他們使用了數據集的一小部分進行測試，只選取了 2 億個訓練樣本，並分別在兩種硬體配置上運行：

配備 1 個 NVIDIA Tesla V100 GPU 和 Xeon Gold 6150 2.70GHz CPU 的伺服器，使用了 PCI Gen 3 介面。

配備 4 個 NVIDIA Tesla V100 GPU 和 IBM Power System AC922 的伺服器，使用了 NVLINK 2.0 介面。

上圖（a）展示的是第一種配置的運行情況。圖中有 S1 和 S2 兩條線。S1 是實際用於訓練的時間，可以看到，每個數據塊的訓練時間大約為 90 毫秒。S2 是用於拷貝數據的時間，從圖中可以看到，拷貝數據需要耗費 318 毫秒。這麼看來，拷貝數據是訓練的瓶頸所在。

上圖（b）展示的是第二種配置的運行情況。可以看到，因為 NVLINK 2.0 介面提供了更大的帶寬，用於拷貝數據的時間降到了 55 毫秒，所以整體速度提升了 3.5 倍。

Snap ML 系統簡析

Snap ML 實現了多個層級的並行，以便在集群的多個節點上分攤工作負載，充分利用加速器單元和計算單元的多核並行能力。

第一層並行

第一層並行是指跨集群多個工作節點運行。訓練數據分布在多個工作節點上，通過網路介面連接到一起。因為有了這種數據並行能力，才能夠打破單台設備的內存限制，訓練大規模的數據集。

第二層並行

在每個工作節點上，可以通過將工作負載分攤給主機和加速器單元來獲得並行能力。不同的工作負載可以同時被執行，工作節點的硬體資源得到了充分利用。

第三層並行

為了更進一步有效利用計算資源，Snap ML 實現了一些特殊的程序，可以充分利用現代 GPU 的並行能力，同時也開發了一些多線程代碼，用於處理 CPU 的工作負載，這就實現了第三個層級的並行。

性能評測

研究人員對 Snap ML 進行了性能評測，讓它運行在不同配置的硬體上，並將它與 sklearn、TensorFlow 和 Apache Spark 進行對比。

評測所使用的機器學習應用程序是互聯網界較為熱門的點擊率預測（Click-Through Rate Prediction，CTR）。CTR 是一種大規模的二元分類應用，基於一系列匿名特徵來預測一個用戶是否會點擊一個廣告。

評測所使用的數據集是由 Criteo 提供的廣告點擊數據集，他們使用前 23 天的數據進行模型訓練，然後使用最後一天的數據進行性能測試。數據集的詳細信息如下表所示：

單節點性能

單節點測試主要是要測出 Snap ML 在單台伺服器上運行的性能，測試內容是訓練一個邏輯回歸分類器，並使用相同的正規化參數（λ = 10），然後比較 Snap ML 與 sklearn 和 TensorFlow 之間的差別。測試使用的是 criteo-kaggle 數據集，它的大小是 11GB，可以塞進單台機器或單個 GPU 的內存里。伺服器是 IBM Power System AC922，配備了 4 個 NVIDIA Tesla V100 GPU，並使用了 NVLINK 2.0 介面，不過在實際運行中只用了其中的一個 GPU。結果如下圖所示。

從圖中可以看出，TensorFlow 使用了大約 500 秒，而 sklearn 只使用了 200 秒左右。這個可能是因為 sklearn 把全部數據集都裝入內存，而 TensorFlow 則是每次從磁碟上讀取一批數據。當然，成績最好的是 Snap ML，只花了大約 20 秒，比前兩個框架快一個數量級，這要歸功於 Snap ML 團隊開發的 GPU 處理程序。

多節點性能

多節點測試主要是要測出 Snap ML 在集群上訓練模型的性能，總共進行了兩次測試。在第一次測試中，本地訓練數據可以完全塞進 GPU 內存中，他們分別訓練了一個線性回歸模型（Linear Regression）和一個支持向量機模型（SVM），並將結果與 Spark MLLib 進行對比。第二次主要測試在本地訓練數據無法完全塞進 GPU 內存的情況下，Snap ML 將會有怎樣的表現。多節點測試使用的是 criteo-1b 數據集，它的大小是 274GB，普通的單台機器無法把它全部塞進內存。

從上圖可以看出，在第一次測試中，在不啟用 GPU 加速的情況下訓練兩種模型，Snap ML 均比 Spark MLLib 快上不止一個數量級：訓練線性回歸模型快 77 倍，訓練支持向量機模型快 24 倍。而在啟用 GPU 加速的情況下，訓練線性回歸模型快 2800 倍，訓練支持向量機模型快 1000 倍。

第二次測試的內容是訓練一個 SVM 分類器。訓練集群包含了 4 個節點，每個節點配備了 2 個 NVIDIA GTX 1080 Ti GPU。這種 GPU 有 11GB 內存，其中 8G 可用於存儲數據。測試使用的二進位數據集 criteo-1b 大小為 98GB，無法全部塞進單台伺服器的 GPU 內存中。為了充分利用每個節點的 GPU，他們結合使用了 DuHL 和 CoCoA。

下圖展示的是啟用 GPU 加速的訓練時間、使用單線程 CPU 進行訓練的時間，以及使用 Snap ML 提供的多線程 CPU 程序進行訓練的時間。

從圖中可以看出，在啟用 GPU 加速的情況下，即使訓練數據無法完全塞進 GPU 內存，因為使用了 DuHL，其訓練速度比使用 CPU 的單線程或多線程都快了一個數量級。

大規模基準測試

這次測試使用了全部的 42 億個訓練樣本，並把 Snap ML 部署在 IBM Power System AC922 伺服器上，每台伺服器配備了 4 個 NVIDIA Tesla V100 GPU。測試內容是訓練一個邏輯回歸模型，並分別與 TensorFlow、LIBLINEAR、Vowpal Wabbit 和 Spark MLLib 已經公布的測試結果進行對比。

從上圖可以看出，Snap ML 的訓練速度比之前最好的成績（TensorFlow）快 46 倍。

如果讀者想要進一步了解 Snap ML 的實現原理和演算法細節，可以閱讀其研究論文：

https://arxiv.org/pdf/1803.06333.pdf

https://www.ibm.com/blogs/research/2018/03/machine-learning-benchmark/

今日薦文

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