基於 Ray 的大規模離線推理

01 大模型離線推理

特點介紹

大數據離線推理

大模型離線推理（Batch 推理）是指在具有數十億至數千億參數的大規模模型上進行分佈式計算推理的過程，具有如下特點：

1. 一次對一批數據進行推理，數據量通常是海量的，所以計算過程通常是離線計算；

2. 推理作業執行過程一般同時包含數據處理及模型推理；

3. 作業規模通常較大，採用分佈式計算，消耗大量計算資源；

4. 相比於在線推理，離線推理對延遲的要求並不高，主要關注吞吐和資源利用率。

關鍵挑戰

GPU Memory Wall

大模型離線推理的關鍵挑戰 — GPU Memory Wall

第一個挑戰是內存的挑戰，機器學習的模型越來越大，尤其是繼 Transformers 類的模型後，模型大小迅猛增長。從上圖中可以看到，過去幾年機器學習領域的模型參數增長非常迅猛，而相比於模型參數的增長，GPU 算力的提升相對較慢，兩者之間就形成了越來越大的 Gap。這就帶來一個問題，在進行推理或者訓練時，GPU 內存可能放不下，需要對模型進行切分。

常見的模型切分方式有上圖左側所列的兩種：

• 按層切分的 Pipeline Parallelism 模式

• 按權重切分的 Tensor Parallelism 模式

按層切分比較簡單，就是將模型的不同層切開，切分成不同的分組，然後放到不同的 GPU 上。比如左上的圖中有兩個 GPU，第一個 GPU 存 L0-L3，第二個 GPU 存 L4-L7。因爲每個層的大小不一樣，所以不一定是平均分配，有的層可能會非常大，獨佔一個 GPU ，小的層就多個擠在一個 GPU 上。

按權重切分就是將模型的同一層，把權重切開放到不同的 GPU 上，比如左下的圖中，將 L0 的一部分權重 A0 放到 GPU 0 上，另外一部分權重 A1 放在 GPU 1 上，在推理的過程中，通過矩陣運算得到最終的結果。除了這兩種方式以外，也有一些更復雜的切分方式，如將這兩種方式進行結合的混合方式，或 Zero 的切分方式。

進行模型切分具有以下幾點優勢：

1. 支持更大模型：可以在現有的硬件基礎上，支持更大模型的離線推理；

2. 降低成本：把現有的模型經過切分之後，放到顯存比較小的卡上，可以降低一部分的成本，那麼更高端的卡就可以出讓給訓練，畢竟訓練會更加消耗資源；

3. 空分複用：目前很多場景會用到空分複用技術，比如英偉達的 Multi-Process Service 技術，即將 GPU 的顯存按照空間切分給不同的進程，能夠提高 GPU 的利用率。但這種情況下，每個進程拿到一部分 GPU 顯存，如果不進行切分，可能要佔據整張卡，所以就是說進行了切分之後，在這種場景下也可以把離線推理運行起來。

分佈式調度

大模型離線推理的關鍵挑戰 — 分佈式調度

第二個挑戰是關於分佈式調度的挑戰。有兩點需求：

第一個是需要支持異構資源，前面說到推理的過程往往同時有數據處理及推理，那麼數據的處理就希望放到 CPU 上進行，從而不佔用 GPU，把 GPU 給到推理使用，所以這就需要框架能夠比較友好地支持異構資源調度。

第二點是對於彈性資源調度的需求，模型經過切分後切成不同的組，在作業的運行過程中，每個組可以理解成一個  Stage，因爲每個組包含的模型的 Layers 是不同的，所以不同 Stage 對於算力的需求也不同，而且在跑一個作業之前，很難預先估計算力需求，就需要不斷地調整參數，才能達到最佳執行效率。所以我們希望計算框架能夠在運行過程中根據計算效率自動對每個 Stage 的算力進行擴縮，使得執行速度快的 Stage 可以自動出讓一些算力給慢的 Stage。

上述兩點需求，目前主流的計算框架，如 Flink 和 Spark，沒有辦法輕易地做到，主要是因爲 Spark 和 Flink 一般綁定了比較固定的批 / 流的計算範式，在調度層面不夠靈活。

性能

性能方面，由於是離線計算作業，我們希望它的吞吐和 GPU 的利用率能夠越高越好。

第一點是數據在 Stage 之間能夠方便且高效的傳輸，應當儘量避免數據落盤帶來的序列化開銷，純內存的傳輸方式是比較好的方式。

第二點是在推理側，應當儘量減少數據 IO 等待，避免 IO 導致 GPU 空閒，最大化提高 GPU 使用率。

第三點是結合資源彈性，釋放掉利用率較低的 GPU，從而提高整理利用率。

案例

案例：Vit + Albert

以下是一個實際的案例，也是一個多模態的例子—— Vit + Albert 雙塔的模型。在這個案例中，我們同時對兩個模型進行切分，一個 GPU 裏面一部分放 Albert 的 Layers，另一部分是 Vit 的 Layers，其中 Embedding 層通常比較大，所以單獨切到一個分組中。作業總共包含了 3 個 Stage，Stage 間傳遞 Image 和文本 Tokerns。因此這 3 個 Stage 所需的計算資源是不同的，即需要彈性分配算力的能力。

02 使用 Ray 構建大模型推理框架

Ray 簡介

Ray 項目是 UC Berkeley 的 RISElab 實驗室在 2017 年前後發起的，RISElab 實驗室的前身是比較著名的 AMP Lab，也就是孵化出了 Spark 引擎的實驗室。該實驗室在更名爲 RISElab 之後，孵化出了 Ray 引擎，Ray 的定位是通用的分佈式編程框架——Python-first。理論上通過 Ray 引擎用戶可以輕鬆地把任何 Python 應用做成分佈式，尤其是機器學習的相關應用，目前 Ray 主攻的一個方向就是機器學習，伯克利的發起者也基於 Ray 創建了創業公司—— Anyscale，目前這個項目在 GitHub 上獲得了兩萬多的關注。在業界，Uber、 OpenAI、螞蟻、字節等公司也都有基於 Ray 的相關應用實踐。

Ray 的架構分爲三層，最下面一層是各種雲基礎設施，也就是說 Ray 幫用戶屏蔽了底層的基礎設施，用戶拉起一個 Ray Cluster 之後就可以立即開始分佈式的編程，不用考慮底層的雲原生或各種各樣的環境；中間層是 Ray Core 層。這一層是 Ray 提供的核心基礎能力，主要是提供了 Low-level 的非常簡潔的分佈式編程 API。基於這套 API，用戶可以非常容易地把現有的 Python 的程序分佈式化。值得注意的是，這一層的 API 是 Low-level，沒有綁定任何的計算範式，非常通用；最上層是 Ray 社區基於 Ray Core 層做的豐富的機器學習庫，這一層的定位是做機器學習 Pipeline。比如，數據加工讀取、模型訓練、超參優化、推理，強化學習等，都可以直接使用這些庫來完成整個的 Pipeline，這也是 Ray 社區目前主攻的一個方向。

更加值得一提的是，據 OpenAI 的公開資料顯示，今年爆火的 ChatGPT，是基於 Ray 進行的包括預訓練、Fine Tune、強化學習等 ChatGPT 的訓練。

Ray 基礎架構

上圖展示的是 Ray Cluster 的基本架構，每一個大框就是一個節點。（這裏的節點是一個虛擬的概念，可以是一個物理機，一個 VM 或一個 Linux 的 Docker。比如在 K8s 上，一個節點就是一個 Pod。）

• Head 節點：是 Ray Cluster 的調度中心，比較核心的組件是 GCS，負責全局存儲、調度、作業、狀態等，Head 節點也有可觀測性 Dashboard。

• Worker 節點：除了 Head 節點之外，其他都是 Worker 節點，承載具體的工作負載。

• Raylet：每個節點上面都有一個守護進程 Raylet，它是一個 Local Scheduler，負責 Task 的調度以及 Worker 的管理。

• Object Store  組件：每個節點上都有 Object Store 組件，負責節點之間 Object 傳輸。在整個 Cluster 中每個節點的 Object Store 組件組成一個全局的分佈式內存。同時，在單個節點上，Object Store 在多進程之間通過共享內存的方式減少 copy。

• Driver：當用戶向 Ray Cluster 上提交一個 Job，或者用 Notebook 連接的時候，Ray 挑選節點來運行 Driver 進行，執行用戶代碼。作業結束後 Driver 銷燬。

• Worker：是 Ray 中 Task 和 Actor 的載體。

此處值得大家關注的是，Ray 爲了提供簡潔的分佈式編程體驗， 在 Raylet 這一層做了非常多的設計，實現過程也比較複雜，感興趣的朋友可以查看相關論文。

Ray 分佈式編程

上圖左側是 Ray Core 的 API 編程：Class 是 Python 的一個類，如果想把它做成分佈式化的話，只需要在類上面加上 @ray.remote 裝飾器，接着創建並調用 Actor 方法，最後通過 ray.get 方法把值取回；因爲 Counter 這個類在遠端的其他節點上，所以我們通過定義一個 Task（Python 函數），使用 Object 進行分佈式的數據傳輸。

右側是使用 Ray 上層的 Library 編程，通過 RayTrain 訓練一個簡單的機器學習模型。使用時需要先定義一個模型，這個過程和直接用 Python 定義模型相同，接着用 RayTrain API 填進去一些 Config 就可以開始訓練。

所以我們看到，這兩種方式一種是 Low-level、一種是 High-level，對於 Ray 來說都是推薦用戶使用的。

基於 Ray 構建大模型推理框架

使用 Ray 構建大模型推理框架 – Ray Datasets

Ray Datasets

在構建大型模型推理框架的選型上，我們選擇了 Ray Datasets。Ray Datasets 提供了豐富的數據源接入方式，兼容目前機器學習領域常用的數據源，並且提供常用的數據處理算子，還支持通用的並行計算，比如在離線的 Bach 推理等。還有一個特點是能夠支持 Pipeline 的執行模式，可以將數據的 Block 劃分爲不同的 Window，大大加速了整個並行計算的執行。總之，Ray Datasets 是一個非常實用的數據處理工具，可以幫助我們更高效地構建大型模型推理框架。

使用 Ray 構建大模型推理框架 v1 — Based on native Ray Dataset Pipeline

因此我們嘗試基於原生的 Ray Datasets Pipeline 構建大模型推理框架。

左邊的僞代碼描述了對應的執行過程，假設將模型按層切分成兩組——ModelLayers1 和 ModelLayers2。調用 Ray Datasets Window API 創建一個 Pipeline，調用 Map Message 在兩個模型分組上進行並行推理。其中 Computer 參數選擇 Actor，表示 Datasets 會在後面爲每一個 Map Batches 的計算過程啓動一個 Actor Pool 進行計算。第三個參數是每個計算 Actor 所需的 GPU 數量， 這個參數會直接作用到背後的 Actor 上，可以看到即使是 Datasets 這類比較高級的庫，它的 API 仍然很容易支持異構資源。

與 Spark 相比，使用 Ray 可以顯著提高執行效率，並且隨着作業規模的擴大，優勢更加明顯。具體在一個簡單的例子中來看，假如我們只有兩個 GPU，模型切了兩組，任務目標是處理三個數據樣本。在使用 Spark 的情況下，需要啓動兩個 Executor 分別加載第一個模型分組的參數並處理 3 個數據樣本，處理後把數據寫到外部存儲中；接下來兩個 Executor 分別再去加載第二個模型分組的參數，然後再分別處理樣本，需要進行跟上一步同樣的處理，最終再將結果寫到外部存儲。由此可見這個過程比較繁瑣，而且對異構資源的支持不太友好。

而使用 Ray 就只需要啓動兩個 Actor，對應 Spark 的兩個 Executor。但是這兩個 Actor 可以分別加載兩個模型分組的參數，兩個 Actor 間的執行過程可以 Pipeline 起來，數據樣本依次經過兩個 Actor。此外也可以非常方便的再增加一個 CPU 上的 Actor 專門做數據的讀取或存儲。框架通過使用 Ray ObjectStore 存儲中間結果數據，純內存存儲避免了序列化的開銷，並且可以顯著提高執行效率，由此可見在此類場景下，使用 Ray 相比於 Spark 可以顯著地提升效率。

使用 Ray 構建大模型推理框架 v1 — Based on native Ray Dataset Pipeline

• 在每個 Actor 進行推理時，數據的 IO 和推理過程並沒有並行起來，導致 GPU 利用率不高；

• Actor Pool 缺乏彈性，在每個 Stage 算力需求不一樣的情況下會浪費資源，需要不斷調整參數以解決這個問題；

• 使用 API 調試參數困難；

• 缺乏容錯和推測執行能力。

爲了解決以上問題，我們開發了第二版推理框架。深入到 Ray Datasets Pipeline 的內部實現中添加了 Streaming 執行語義。各個 Stage 通過 Queue 前後連接起來，Queue 中傳遞的是 Ray Object Reference 而不是實際數據，實際數據在 Actor 側。相當於我們寫程序時函數之間傳遞指針數組而不是實際數據。

第二版推理框架和第一版不同，每一個 Stage 背後是一個穩定的 Actor Pool，從一開始被創建之後就不會釋放。在運行的過程中，該 Stage 就從它的 Input Queue 中讀取 Object Reference，讀到數據後在自己的 Actor Pool 中選擇一個 Actor 來處理數據。因爲 Actor Pool 是自定義的，可以實現彈性能力，使負載重的 Stage 的 Actor Pool 會主動嘗試申請更多的資源來增加自己的並行度，而負載輕的 Stage 的 Actor Pool 會逐漸空閒下來，並最終釋放掉一些 Actor，從而出讓資源給需要資源更多的 Stage。當然，這個也需要配合一定的調度策略，也就是 Stage 在分發數據的時候如何選擇一個 Actor。我們現在使用的是 Most Recently Used 的策略，忙的 Actor 就讓他更忙，這樣空閒的 Actor 就可以容易空閒下來釋放掉。在 Actor 側，利用 Actor 內多線程實現 IO 和推理計算並行，提高了 GPU 的利用率。

需要注意的是，Stage 之間 Queue 的長度是有限的，可以避免上游的 Stage 產生過多的數據導致作業 OOM，相當於流計算中反壓的作用。第二版的僞代碼和第一版並沒有太多的不同，因此業務不需要花費很大的精力進行改造。

使用 Ray 構建大模型推理框架 v2 — v2.3 社區合作

在開發第二版的同事，注意到 Ray 開源社區也在考慮類似的問題。社區提出了一個官方 REP，其中列出的問題與我們的目標非常相似，尤其是在提高 GPU 利用率和解決 Ray Datasets API 參數配置困難這兩個方面。社區提出了新的架構，在 API 層下面拆分出了 Operator 和 Executor，增加了更多的靈活性和擴展性。後面，我們也會和社區展開合作，將我們的實現做爲新架構下的一種 Executor。

03 Ray 雲原生部署實踐

在部署 Ray 時，我們使用了開源社區完整的解決方案 Kuberay 項目。前面提到每個 Ray Cluster 由 Head 節點和 Worker 節點組成，每個節點是一份計算資源，可以是物理機、Docker 等等，在 K8s 上即爲一個 Pod。啓動 Ray Cluster 時，使用 Kuberay 的 Operator 來管理整個生命週期，包括創建和銷燬 Cluster 等等。目前，這個項目也比較活躍，字節、微軟、螞蟻等公司都有參與研發和使用。

在字節內部，用戶可以通過內部的平臺使用 Ray，通過提交 Job 或使用 Notebook 進行交互式編程。平臺通過 Kuberay 提供的 YAML 和 Restful API 這兩種方式進行操作。Kuberay 同時也支持自動擴展和水平擴展。Ray Cluster 在內部用於收集負載的 Metrics，並根據 Metrics 決定是否擴充更多的資源，如果需要則觸發 Kuberay 拉起新的 Pod 或刪除閒置的 Pod。

最後總結一下，我們今天討論了大模型離線推理以及其中關鍵的挑戰，並介紹瞭如何使用 Ray 構建大模型推理框架。未來，我們將繼續加強與社區的合作，優化我們的平臺，並挖掘更多的 Ray 上的應用場景。

本文由 Readfog 進行 AMP 轉碼，版權歸原作者所有。
來源：https://mp.weixin.qq.com/s/2-jWtYcO0CVnttRrJOYcnA