大模型離線推理

  特點(diǎn)介紹

大數(shù)據(jù)離線推理

大模型離線推理（Batch 推理）是指在具有數(shù)十億至數(shù)千億參數(shù)的大規(guī)模模型上進(jìn)行分布式計算推理的過程，具有如下特點(diǎn)：

1. 一次對一批數(shù)據(jù)進(jìn)行推理，數(shù)據(jù)量通常是海量的，所以計算過程通常是離線計算；
2. 推理作業(yè)執(zhí)行過程一般同時包含數(shù)據(jù)處理及模型推理；
3. 作業(yè)規(guī)模通常較大，采用分布式計算，消耗大量計算資源；
4. 相比于在線推理，離線推理對延遲的要求并不高，主要關(guān)注吞吐和資源利用率。

  關(guān)鍵挑戰(zhàn)

GPU Memory Wall

圖片

大模型離線推理的關(guān)鍵挑戰(zhàn) — GPU Memory Wall

第一個挑戰(zhàn)是內(nèi)存的挑戰(zhàn)，機(jī)器學(xué)習(xí)的模型越來越大，尤其是繼 Transformers 類的模型后，模型大小迅猛增長。從上圖中可以看到，過去幾年機(jī)器學(xué)習(xí)領(lǐng)域的模型參數(shù)增長非常迅猛，而相比于模型參數(shù)的增長，GPU 算力的提升相對較慢，兩者之間就形成了越來越大的 Gap。這就帶來一個問題，在進(jìn)行推理或者訓(xùn)練時，GPU 內(nèi)存可能放不下，需要對模型進(jìn)行切分。

圖片

模型切分

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

• 按層切分的 Pipeline Parallelism 模式
• 按權(quán)重切分的 Tensor Parallelism 模式

按層切分比較簡單，就是將模型的不同層切開，切分成不同的分組，然后放到不同的 GPU 上。比如左上的圖中有兩個GPU，第一個 GPU 存 L0-L3，第二個 GPU 存 L4-L7。因?yàn)槊總€層的大小不一樣，所以不一定是平均分配，有的層可能會非常大，獨(dú)占一個 GPU ，小的層就多個擠在一個 GPU 上。

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

進(jìn)行模型切分具有以下幾點(diǎn)優(yōu)勢：

1. 支持更大模型：可以在現(xiàn)有的硬件基礎(chǔ)上，支持更大模型的離線推理；
2. 降低成本：把現(xiàn)有的模型經(jīng)過切分之后，放到顯存比較小的卡上，可以降低一部分的成本，那么更高端的卡就可以出讓給訓(xùn)練，畢竟訓(xùn)練會更加消耗資源；
3. 空分復(fù)用：目前很多場景會用到空分復(fù)用技術(shù)，比如英偉達(dá)的 Multi-Process Service 技術(shù)，即將 GPU 的顯存按照空間切分給不同的進(jìn)程，能夠提高 GPU 的利用率。但這種情況下，每個進(jìn)程拿到一部分 GPU 顯存，如果不進(jìn)行切分，可能要占據(jù)整張卡，所以就是說進(jìn)行了切分之后，在這種場景下也可以把離線推理運(yùn)行起來。

分布式調(diào)度

圖片

大模型離線推理的關(guān)鍵挑戰(zhàn) — 分布式調(diào)度

第二個挑戰(zhàn)是關(guān)于分布式調(diào)度的挑戰(zhàn)。有兩點(diǎn)需求：

第一個是需要支持異構(gòu)資源，前面說到推理的過程往往同時有數(shù)據(jù)處理及推理，那么數(shù)據(jù)的處理就希望放到 CPU 上進(jìn)行，從而不占用 GPU，把 GPU 給到推理使用，所以這就需要框架能夠比較友好地支持異構(gòu)資源調(diào)度。

第二點(diǎn)是對于彈性資源調(diào)度的需求，模型經(jīng)過切分后切成不同的組，在作業(yè)的運(yùn)行過程中，每個組可以理解成一個  Stage，因?yàn)槊總€組包含的模型的 Layers 是不同的，所以不同 Stage 對于算力的需求也不同，而且在跑一個作業(yè)之前，很難預(yù)先估計算力需求，就需要不斷地調(diào)整參數(shù)，才能達(dá)到最佳執(zhí)行效率。所以我們希望計算框架能夠在運(yùn)行過程中根據(jù)計算效率自動對每個 Stage 的算力進(jìn)行擴(kuò)縮，使得執(zhí)行速度快的 Stage 可以自動出讓一些算力給慢的 Stage。

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

性能

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

第一點(diǎn)是數(shù)據(jù)在 Stage 之間能夠方便且高效的傳輸，應(yīng)當(dāng)盡量避免數(shù)據(jù)落盤帶來的序列化開銷，純內(nèi)存的傳輸方式是比較好的方式。

第二點(diǎn)是在推理側(cè)，應(yīng)當(dāng)盡量減少數(shù)據(jù) IO 等待，避免 IO 導(dǎo)致 GPU 空閑，最大化提高 GPU 使用率。

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

案例

圖片

案例：Vit + Albert

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

02

使用 Ray 構(gòu)建大模型推理框架

  Ray 簡介

圖片

Ray 簡介

Ray 項(xiàng)目是 UC Berkeley 的 RISElab 實(shí)驗(yàn)室在 2017 年前后發(fā)起的，RISElab 實(shí)驗(yàn)室的前身是比較著名的 AMP Lab，也就是孵化出了 Spark 引擎的實(shí)驗(yàn)室。該實(shí)驗(yàn)室在更名為 RISElab 之后，孵化出了 Ray 引擎，Ray 的定位是通用的分布式編程框架——Python-first。理論上通過 Ray 引擎用戶可以輕松地把任何 Python 應(yīng)用做成分布式，尤其是機(jī)器學(xué)習(xí)的相關(guān)應(yīng)用，目前 Ray 主攻的一個方向就是機(jī)器學(xué)習(xí)，伯克利的發(fā)起者也基于 Ray 創(chuàng)建了創(chuàng)業(yè)公司—— Anyscale，目前這個項(xiàng)目在 GitHub 上獲得了兩萬多的關(guān)注。在業(yè)界，Uber、 OpenAI、螞蟻、字節(jié)等公司也都有基于 Ray 的相關(guān)應(yīng)用實(shí)踐。

Ray 的架構(gòu)分為三層，最下面一層是各種云基礎(chǔ)設(shè)施，也就是說 Ray 幫用戶屏蔽了底層的基礎(chǔ)設(shè)施，用戶拉起一個 Ray Cluster之后就可以立即開始分布式的編程，不用考慮底層的云原生或各種各樣的環(huán)境；中間層是 Ray Core 層。這一層是 Ray 提供的核心基礎(chǔ)能力，主要是提供了 Low-level 的非常簡潔的分布式編程 API?；谶@套 API，用戶可以非常容易地把現(xiàn)有的 Python 的程序分布式化。值得注意的是，這一層的 API 是 Low-level，沒有綁定任何的計算范式，非常通用；最上層是 Ray 社區(qū)基于 Ray Core 層做的豐富的機(jī)器學(xué)習(xí)庫，這一層的定位是做機(jī)器學(xué)習(xí) Pipeline。比如，數(shù)據(jù)加工讀取、模型訓(xùn)練、超參優(yōu)化、推理，強(qiáng)化學(xué)習(xí)等，都可以直接使用這些庫來完成整個的 Pipeline，這也是 Ray 社區(qū)目前主攻的一個方向。

更加值得一提的是，據(jù) OpenAI 的公開資料顯示，今年爆火的 ChatGPT，是基于 Ray 進(jìn)行的包括預(yù)訓(xùn)練、Fine Tune、強(qiáng)化學(xué)習(xí)等 ChatGPT 的訓(xùn)練。

  Ray 基礎(chǔ)架構(gòu)

圖片

Ray 基礎(chǔ)架構(gòu)

上圖展示的是 Ray Cluster 的基本架構(gòu)，每一個大框就是一個節(jié)點(diǎn)。（這里的節(jié)點(diǎn)是一個虛擬的概念，可以是一個物理機(jī)，一個 VM 或一個 Linux 的 Docker。比如在 K8s 上，一個節(jié)點(diǎn)就是一個 Pod。）

• Head 節(jié)點(diǎn)：是 Ray Cluster 的調(diào)度中心，比較核心的組件是 GCS，負(fù)責(zé)全局存儲、調(diào)度、作業(yè)、狀態(tài)等，Head節(jié)點(diǎn)也有可觀測性 Dashboard。
• Worker 節(jié)點(diǎn)：除了 Head 節(jié)點(diǎn)之外，其他都是 Worker 節(jié)點(diǎn)，承載具體的工作負(fù)載。

• Raylet：每個節(jié)點(diǎn)上面都有一個守護(hù)進(jìn)程 Raylet，它是一個 Local Scheduler，負(fù)責(zé) Task 的調(diào)度以及 Worker 的管理。
• Object Store  組件：每個節(jié)點(diǎn)上都有 Object Store 組件，負(fù)責(zé)節(jié)點(diǎn)之間 Object 傳輸。在整個 Cluster 中每個節(jié)點(diǎn)的 Object Store 組件組成一個全局的分布式內(nèi)存。同時，在單個節(jié)點(diǎn)上，Object Store 在多進(jìn)程之間通過共享內(nèi)存的方式減少 copy。

• Driver：當(dāng)用戶向 Ray Cluster 上提交一個 Job，或者用 Notebook 連接的時候，Ray挑選節(jié)點(diǎn)來運(yùn)行 Driver 進(jìn)行，執(zhí)行用戶代碼。作業(yè)結(jié)束后 Driver 銷毀。
• Worker：是 Ray 中 Task 和 Actor 的載體。

此處值得大家關(guān)注的是，Ray 為了提供簡潔的分布式編程體驗(yàn)， 在 Raylet 這一層做了非常多的設(shè)計，實(shí)現(xiàn)過程也比較復(fù)雜，感興趣的朋友可以查看相關(guān)論文。

  Ray 分布式編程

圖片

Ray 分布式編程

上圖左側(cè)是 Ray Core 的 API 編程：Class 是 Python 的一個類，如果想把它做成分布式化的話，只需要在類上面加上 @ray.remote 裝飾器，接著創(chuàng)建并調(diào)用 Actor 方法，最后通過 ray.get 方法把值取回；因?yàn)?Counter 這個類在遠(yuǎn)端的其他節(jié)點(diǎn)上，所以我們通過定義一個 Task（Python 函數(shù)），使用 Object 進(jìn)行分布式的數(shù)據(jù)傳輸。

右側(cè)是使用 Ray 上層的 Library 編程，通過 RayTrain 訓(xùn)練一個簡單的機(jī)器學(xué)習(xí)模型。使用時需要先定義一個模型，這個過程和直接用 Python 定義模型相同，接著用 RayTrain API 填進(jìn)去一些 Config 就可以開始訓(xùn)練。

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

  基于 Ray 構(gòu)建大模型推理框架

圖片

使用 Ray 構(gòu)建大模型推理框架 – Ray Datasets

Ray Datasets

在構(gòu)建大型模型推理框架的選型上，我們選擇了 Ray Datasets。Ray Datasets 提供了豐富的數(shù)據(jù)源接入方式，兼容目前機(jī)器學(xué)習(xí)領(lǐng)域常用的數(shù)據(jù)源，并且提供常用的數(shù)據(jù)處理算子，還支持通用的并行計算，比如在離線的 Bach 推理等。還有一個特點(diǎn)是能夠支持 Pipeline 的執(zhí)行模式，可以將數(shù)據(jù)的 Block 劃分為不同的 Window，大大加速了整個并行計算的執(zhí)行?？傊?，Ray Datasets 是一個非常實(shí)用的數(shù)據(jù)處理工具，可以幫助我們更高效地構(gòu)建大型模型推理框架。

圖片

使用 Ray 構(gòu)建大模型推理框架 v1 — Based on native Ray Dataset Pipeline

因此我們嘗試基于原生的 Ray Datasets Pipeline 構(gòu)建大模型推理框架。

左邊的偽代碼描述了對應(yīng)的執(zhí)行過程，假設(shè)將模型按層切分成兩組——ModelLayers1 和 ModelLayers2。調(diào)用 Ray Datasets Window API 創(chuàng)建一個 Pipeline，調(diào)用 Map Message 在兩個模型分組上進(jìn)行并行推理。其中 Computer 參數(shù)選擇 Actor，表示 Datasets 會在后面為每一個 Map Batches 的計算過程啟動一個 Actor Pool 進(jìn)行計算。第三個參數(shù)是每個計算 Actor 所需的 GPU 數(shù)量， 這個參數(shù)會直接作用到背后的 Actor 上，可以看到即使是 Datasets 這類比較高級的庫，它的 API 仍然很容易支持異構(gòu)資源。

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

而使用 Ray 就只需要啟動兩個 Actor，對應(yīng) Spark 的兩個 Executor。但是這兩個 Actor 可以分別加載兩個模型分組的參數(shù)，兩個 Actor 間的執(zhí)行過程可以 Pipeline 起來，數(shù)據(jù)樣本依次經(jīng)過兩個 Actor。此外也可以非常方便的再增加一個 CPU 上的 Actor 專門做數(shù)據(jù)的讀取或存儲?？蚣芡ㄟ^使用 Ray ObjectStore 存儲中間結(jié)果數(shù)據(jù)，純內(nèi)存存儲避免了序列化的開銷，并且可以顯著提高執(zhí)行效率，由此可見在此類場景下，使用 Ray 相比于 Spark 可以顯著地提升效率。

圖片

使用 Ray 構(gòu)建大模型推理框架 v1 — Based on native Ray Dataset Pipeline

上面介紹的是第一版的大模型推理框架，它已經(jīng)能夠很好地解決異構(gòu)資源調(diào)度的問題，并且在理論上，它的執(zhí)行效率也超過了同類的計算引擎，但是我們發(fā)現(xiàn)仍然存在一些問題：

• 每個 Window 都要創(chuàng)建和銷毀 Actor Pool，而每次拉起 Actor 加載模型的開銷太大；
• 在每個 Actor 進(jìn)行推理時，數(shù)據(jù)的 IO 和推理過程并沒有并行起來，導(dǎo)致 GPU 利用率不高；
• Actor Pool 缺乏彈性，在每個 Stage 算力需求不一樣的情況下會浪費(fèi)資源，需要不斷調(diào)整參數(shù)以解決這個問題；
• 使用 API 調(diào)試參數(shù)困難；
• 缺乏容錯和推測執(zhí)行能力。

圖片

使用 Ray 構(gòu)建大模型推理框架 v2 — Streaming execution semantics in Ray Dataset Pipeline

為了解決以上問題，我們開發(fā)了第二版推理框架。深入到 Ray Datasets Pipeline 的內(nèi)部實(shí)現(xiàn)中添加了 Streaming 執(zhí)行語義。各個 Stage 通過 Queue 前后連接起來，Queue 中傳遞的是 Ray Object Reference 而不是實(shí)際數(shù)據(jù)，實(shí)際數(shù)據(jù)在 Actor 側(cè)。相當(dāng)于我們寫程序時函數(shù)之間傳遞指針數(shù)組而不是實(shí)際數(shù)據(jù)。

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

需要注意的是，Stage 之間 Queue 的長度是有限的，可以避免上游的 Stage 產(chǎn)生過多的數(shù)據(jù)導(dǎo)致作業(yè) OOM，相當(dāng)于流計算中反壓的作用。第二版的偽代碼和第一版并沒有太多的不同，因此業(yè)務(wù)不需要花費(fèi)很大的精力進(jìn)行改造。

圖片

使用 Ray 構(gòu)建大模型推理框架 v2 — v2.3 社區(qū)合作

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

03

Ray 云原生部署實(shí)踐

圖片

Ray 云原生部署 — KubeRay

在部署 Ray 時，我們使用了開源社區(qū)完整的解決方案 Kuberay 項(xiàng)目。前面提到每個 Ray Cluster 由 Head 節(jié)點(diǎn)和 Worker 節(jié)點(diǎn)組成，每個節(jié)點(diǎn)是一份計算資源，可以是物理機(jī)、Docker 等等，在 K8s 上即為一個 Pod。啟動 Ray Cluster 時，使用 Kuberay 的 Operator 來管理整個生命周期，包括創(chuàng)建和銷毀 Cluster 等等。目前，這個項(xiàng)目也比較活躍，字節(jié)、微軟、螞蟻等公司都有參與研發(fā)和使用。

圖片

Ray 云原生部署 — KubeRay

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

最后總結(jié)一下，我們今天討論了大模型離線推理以及其中關(guān)鍵的挑戰(zhàn)，并介紹了如何使用 Ray 構(gòu)建大模型推理框架。未來，我們將繼續(xù)加強(qiáng)與社區(qū)的合作，優(yōu)化我們的平臺，并挖掘更多的 Ray 上的應(yīng)用場景。